Coverage for src/bartz/mcmcstep/_state.py: 98%
549 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 11:03 +0000
1# bartz/src/bartz/mcmcstep/_state.py
2#
3# Copyright (c) 2024-2026, The Bartz Contributors
4#
5# This file is part of bartz.
6#
7# Permission is hereby granted, free of charge, to any person obtaining a copy
8# of this software and associated documentation files (the "Software"), to deal
9# in the Software without restriction, including without limitation the rights
10# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11# copies of the Software, and to permit persons to whom the Software is
12# furnished to do so, subject to the following conditions:
13#
14# The above copyright notice and this permission notice shall be included in all
15# copies or substantial portions of the Software.
16#
17# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23# SOFTWARE.
25"""Module defining the BART MCMC state and initialization."""
27from collections.abc import Callable, Sequence
28from dataclasses import dataclass, replace
29from enum import Enum
30from functools import partial, wraps
31from typing import Literal, NamedTuple, TypedDict, TypeVar, cast
33import jax
34import numpy
35from equinox import error_if, filter_jit
36from jax import NamedSharding, device_put, lax, make_mesh, random, shard_map, tree, vmap
37from jax import numpy as jnp
38from jax.scipy.linalg import solve_triangular
39from jax.sharding import AxisType, Mesh, PartitionSpec
40from jax.typing import DTypeLike
41from jaxtyping import (
42 Array,
43 Bool,
44 Float,
45 Float32,
46 Int32,
47 Integer,
48 Key,
49 PyTree,
50 Shaped,
51 UInt,
52 UInt32,
53)
54from numpy import ndarray
56from bartz._jaxext import (
57 Module,
58 field,
59 float32_matmuls,
60 jaxtyping_disabled,
61 jit,
62 minimal_unsigned_dtype,
63)
64from bartz.grove import tree_depths
65from bartz.mcmcstep._axes import CHAIN_AXIS, chain_vmap_axes, data_vmap_axes
66from bartz.mcmcstep._lazy import (
67 _is_lazy_or_none,
68 _lazy_from_array,
69 _LazyArray,
70 _wrap_chain,
71 add_dummy_axis,
72 lazy,
73)
74from bartz.mcmcstep._reduction import (
75 AutoBatchedReduction,
76 AutoOneHotReduction,
77 ReductionConfig,
78)
80ArrayLike = Array | ndarray
82FloatLike = float | Float[ArrayLike, '']
85class OutcomeType(Enum):
86 """Likelihood types for each outcome component in the regression."""
88 continuous = 'continuous'
89 """Continuous outcome with Normal error."""
91 binary = 'binary'
92 """Binary outcome in {0, 1} with probit link."""
95T = TypeVar('T')
98class Wishart(Module):
99 """A precision matrix with a Wishart prior, bundled with its current value.
101 Represents a random precision (inverse covariance) ``value`` drawn from a
102 Wishart prior with degrees of freedom `nu` and rate matrix `rate`. The
103 univariate case (``k = 1``) is the Gamma special case; the relationship to
104 the inverse-gamma prior on the variance is ``alpha = nu / 2``,
105 ``beta = rate / 2``. The prior mean of the precision is ``nu * rate^-1``.
107 Set `nu` and `rate` to `None` to represent a precision held fixed at `value`
108 with no prior (e.g. the identity in binary regression).
109 """
111 nu: Float32[Array, ''] | None
112 """Degrees of freedom of the Wishart prior, or `None` if there is no prior."""
114 rate: Float32[Array, ''] | Float32[Array, 'k k'] | None
115 """The rate matrix of the Wishart prior (scalar for univariate), or `None`
116 if there is no prior. Equal to the inverse-gamma ``scale`` in the
117 univariate case."""
119 value: Float32[Array, '*chains k k'] | Float32[Array, '*chains'] = field(
120 chains=CHAIN_AXIS
121 )
122 """The precision matrix (scalar for univariate)."""
124 def __init__(
125 self,
126 nu: FloatLike | None,
127 rate: FloatLike | Float[ArrayLike, 'k k'] | None,
128 value: FloatLike
129 | Float[ArrayLike, '*chains k k']
130 | Float[ArrayLike, '*chains'],
131 ) -> None:
132 # `init` passes a deferred `_LazyArray` (cast to `Array`) for `value` to
133 # route it through sharding.
134 assert (nu is None) == (rate is None), 'set both or neither of nu and rate'
135 self.nu = None if nu is None else jnp.asarray(nu, jnp.float32)
136 self.rate = None if rate is None else jnp.asarray(rate, jnp.float32)
137 if isinstance(value, _LazyArray):
138 self.value = cast(Array, value)
139 else:
140 self.value = jnp.asarray(value, jnp.float32)
143class DiagWishart(Wishart):
144 """A diagonal precision matrix with independent chi-square diagonal entries.
146 Despite the name this is not a Wishart restricted to diagonal matrices, but
147 a convenience type: a diagonal precision whose entries are mutually
148 independent, each with its own Gamma (scaled chi-square) prior. Only the
149 multivariate (matrix) case is supported.
151 A component with `rate` 0 has no prior; its precision is held fixed at its
152 `value` (1 for the binary components of a mixed regression).
154 Used for mixed binary-continuous regression and for continuous multivariate
155 regression with per-datapoint missingness.
156 """
158 def __init__(
159 self,
160 nu: FloatLike | None,
161 rate: FloatLike | Float[ArrayLike, 'k k'] | None,
162 value: FloatLike
163 | Float[ArrayLike, '*chains k k']
164 | Float[ArrayLike, '*chains'],
165 ) -> None:
166 # explicit (delegating) init so the static checker uses this signature
167 # instead of synthesizing a stricter one from the inherited fields
168 assert rate is None or jnp.ndim(rate) == 2, (
169 'DiagWishart supports only the multivariate (matrix) case'
170 )
171 super().__init__(nu, rate, value)
174class Forest(Module):
175 """Represents the MCMC state of a sum of trees."""
177 # Implementation note related to runtime type checking: Heap-array fields
178 # follow the `bartz.grove.TreesTrace` convention: the union-free integer
179 # trees are declared before `leaf_tree` and carry the bindable
180 # `half_tree_size` axis, while `leaf_tree` (and `p_nonterminal`) are
181 # checked against `2*half_tree_size`. Declaring a union-free `*chains`
182 # field first binds the variadic chain axis (plus `num_trees` and
183 # `half_tree_size`) before `leaf_tree`'s `... | ... k ...` union is
184 # evaluated, so the runtime typechecker can't mis-bind `*chains` against
185 # the `k` axis of a multivariate forest (the layouts are otherwise
186 # rank-ambiguous). No dummy anchor field needed.
188 var_tree: UInt[Array, '*chains num_trees half_tree_size'] = field(chains=CHAIN_AXIS)
189 """Variables/predictors/axes of decision rules."""
191 split_tree: UInt[Array, '*chains num_trees half_tree_size'] = field(
192 chains=CHAIN_AXIS
193 )
194 """Cutpoints/boundaries of decision rules."""
196 affluence_tree: Bool[Array, '*chains num_trees half_tree_size'] = field(
197 chains=CHAIN_AXIS
198 )
199 """Marks leaves that can be grown."""
201 leaf_tree: (
202 Float[Array, '*chains num_trees 2*half_tree_size']
203 | Float[Array, '*chains num_trees k 2*half_tree_size']
204 ) = field(chains=CHAIN_AXIS)
205 """The leaf values, in units of `leaf_unit`. The function computed by the
206 forest is ``offset + leaf_unit * (sum of leaves)``."""
208 leaf_unit: Float32[Array, ''] | Float32[Array, ' k']
209 """The storage unit of the leaves. Keeps the stored values O(1) whatever
210 the data units, so they do not under/overflow narrow `leaf_tree` dtypes.
211 Set to the marginal prior standard deviation of a leaf, rounded to a power
212 of two so converting to and from data units is exact."""
214 offset: Float32[Array, ''] | Float32[Array, ' k']
215 """Constant shift added to the scaled sum of trees, see `leaf_tree`."""
217 grow_prop_count: Int32[Array, '*chains'] = field(chains=CHAIN_AXIS)
218 """The number of grow proposals made during the last MCMC step."""
220 prune_prop_count: Int32[Array, '*chains'] = field(chains=CHAIN_AXIS)
221 """The number of prune proposals made during the last MCMC step."""
223 grow_acc_count: Int32[Array, '*chains'] = field(chains=CHAIN_AXIS)
224 """The number of grow moves accepted during the last MCMC step."""
226 prune_acc_count: Int32[Array, '*chains'] = field(chains=CHAIN_AXIS)
227 """The number of prune moves accepted during the last MCMC step."""
229 max_split: UInt[Array, ' p']
230 """The maximum split index for each predictor."""
232 blocked_vars: UInt[Array, ' q'] | None
233 """Indices of not to be used variables/predictors. This shall include at
234 least all the `i` that yield ``max_split[i] == 0``, otherwise behavior is
235 undefined."""
237 p_nonterminal: Float32[Array, ' 2*half_tree_size']
238 """The prior probability of each node being nonterminal (conditional on its
239 ancestors leaving at least one available decision rule). Includes the nodes
240 at maximum depth which shall be set to 0."""
242 p_propose_grow: Float32[Array, ' half_tree_size']
243 """The unnormalized probability of picking a leaf for a grow proposal."""
245 leaf_indices: UInt[Array, '*chains num_trees n'] = field(chains=CHAIN_AXIS, data=-1)
246 """The index of the leaf each datapoint falls into, for each tree, in the
247 largest version of the tree compatible with the last moves.
249 A pending prune (accepted prune or rejected grow, marked per-tree by
250 `to_prune`) is not yet applied to the indices; `step` folds it in at the
251 beginning of the next iteration. Evaluating the trees at these indices is
252 correct anyway because `leaf_tree` mirrors the value of a pruned node onto
253 its dangling children."""
255 to_prune: Bool[Array, '*chains num_trees'] = field(chains=CHAIN_AXIS)
256 """Whether the last move on each tree ended in a prune (accepted prune or
257 rejected grow) whose application to `leaf_indices` is still pending."""
259 move_node: Int32[Array, '*chains num_trees'] = field(chains=CHAIN_AXIS)
260 """The node the last move on each tree operated on (the leaf to grow or
261 the node to prune). Meaningful only where `to_prune` is set."""
263 count_tree: UInt32[Array, '*chains num_trees 2*half_tree_size'] | None = field(
264 chains=CHAIN_AXIS
265 )
266 """The number of datapoints per leaf. Valid at the leaves and at the nodes
267 involved in the latest moves, dirty elsewhere. `None` if there are
268 per-datapoint error scales and no minimum-points-per-node constraints,
269 which makes the counts unused."""
271 prec_tree: (
272 Float32[Array, '*chains num_trees 2*half_tree_size']
273 | Float32[Array, '*chains num_trees k k 2*half_tree_size']
274 | None
275 ) = field(chains=CHAIN_AXIS)
276 """The sum of `State.prec_scale` over the datapoints in each leaf, in the
277 same units; valid/dirty like `count_tree`. `None` if there are no
278 per-datapoint error scales, in which case `count_tree` takes its place."""
280 min_points_per_decision_node: Int32[Array, ''] | None
281 """The minimum number of data points in a decision node."""
283 min_points_per_leaf: Int32[Array, ''] | None
284 """The minimum number of data points in a leaf node."""
286 log_trans_prior: Float32[Array, '*chains num_trees'] | None = field(
287 chains=CHAIN_AXIS
288 )
289 """The log transition and prior Metropolis-Hastings ratio for the proposed
290 move on each tree."""
292 log_likelihood: Float32[Array, '*chains num_trees'] | None = field(
293 chains=CHAIN_AXIS
294 )
295 """The log likelihood ratio."""
297 leaf_prior_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'] | None
298 """The prior precision matrix of a leaf, conditional on the tree structure
299 (a scalar inverse variance for univariate). The prior mean of a leaf is
300 zero; the prior covariance of the sum of trees is ``num_trees *
301 leaf_prior_cov_inv^-1``."""
303 log_s: Float32[Array, '*chains p'] | None = field(chains=CHAIN_AXIS)
304 """The logarithm of the prior probability for choosing a variable to split
305 along in a decision rule. Not normalized. Variables that do not have any
306 available decision rule at a given node are masked away. If `None`, use a
307 uniform distribution."""
309 theta: Float32[Array, '*chains'] | None = field(chains=CHAIN_AXIS)
310 """The concentration parameter for the Dirichlet prior on the variable
311 distribution `s`. If not set, `log_s` is left constant."""
313 a: Float32[Array, ''] | None
314 """Parameter of the prior on `theta`. If not set, `theta` is left constant."""
316 b: Float32[Array, ''] | None
317 """Parameter of the prior on `theta`. If not set, `theta` is left constant."""
319 rho: Float32[Array, ''] | None
320 """Parameter of the prior on `theta`. If not set, `theta` is left constant."""
322 @property
323 def has_chains(self) -> bool:
324 """Whether non-constant attributes in the forest carry an explicit chain axis."""
325 return self.var_tree.ndim > 2
328class StepConfig(Module):
329 """Options for the MCMC step."""
331 steps_done: Int32[Array, '']
332 """The number of MCMC steps completed so far."""
334 sparse_on_at: Int32[Array, ''] | None
335 """After how many steps to turn on variable selection. If `None`, variable
336 selection is disabled."""
338 resid_reduction_config: ReductionConfig
339 """How to sum the residuals in each leaf."""
341 count_reduction_config: ReductionConfig
342 """How to count the datapoints in each leaf."""
344 prec_reduction_config: ReductionConfig
345 """How to sum the likelihood precisions in each leaf."""
347 prec_count_num_trees: int | None = field(static=True)
348 """Batch size for processing trees to compute count and prec trees."""
350 sequential_unroll: int | bool = field(static=True)
351 """How much to unroll the sequential accept/reject loop over trees in
352 `step`. See the ``unroll`` argument of `jax.lax.scan`."""
354 augment: bool = field(static=True)
355 """Whether to account exactly, via data augmentation, for the decision
356 rules forbidden by the ancestors of each node when updating
357 `Forest.log_s`."""
359 mesh: Mesh | None = field(static=True)
360 """The mesh used to shard data and computation across multiple devices."""
362 leaf_quantization: Int32[Array, ''] | None = None
363 """If set, quantize the leaves to multiples of ``eps(resid dtype) * 2 **
364 leaf_quantization`` in `State.resid_eff_scale` units, which makes the
365 running updates of `State.resid` mostly exact, (almost) stopping their
366 random-walk rounding drift, assuming ``|resid| < 2 ** (leaf_quantization +
367 1)`` holds (in the same units) for most datapoints most of the time.
368 Intended mostly for use with float16 residuals. Sensible settings are 0 and
369 1, with 0 leaving some drift, 1 practically no drift, and no setting above
370 1 justifying the reduced accuracy. With enough datapoints or trees the
371 MCMC breaks down because the sampled leaf variation becomes smaller than
372 the quantum, so this setting can not be used liberally."""
374 @property
375 def data_sharded(self) -> bool:
376 """Whether the data axis is sharded across devices."""
377 return self.mesh is not None and 'data' in self.mesh.axis_names
380class State(Module):
381 """Represents the MCMC state of BART."""
383 _chain_anchor: Float32[Array, '*chains'] = field(chains=CHAIN_AXIS)
384 """Unused per-chain scalar, declared first as a runtime-typechecker anchor.
385 Its single (union-free) ``*chains`` annotation binds the variadic chain
386 axis before the ``... | ... k ...`` unions of `z`/`resid` (z over the
387 binary-outcome ``kb`` axis) are checked; otherwise those can mis-bind
388 ``*chains`` against the outcome axis for a multivariate-without-chains
389 state (the layouts are rank-ambiguous). Unlike `Forest`, `State` has no
390 genuine union-free chain field to reorder into this slot, so a dummy one is
391 carried."""
393 X: UInt[Array, 'p n'] = field(data=-1)
394 """The predictors."""
396 y: Float32[Array, ' n'] | Float32[Array, 'k n'] = field(data=-1)
397 """The response, in data units. Binary components are stored as 0/1.
398 Missing values are replaced by `Forest.offset`."""
400 z: None | Float32[Array, '*chains n'] | Float32[Array, '*chains kb n'] = field(
401 chains=CHAIN_AXIS, data=-1
402 )
403 """The latent outcomes for binary regression. `None` in continuous
404 regression. In the mixed binary-continuous case, only the binary outcome
405 components are stored."""
407 binary_indices: None | Int32[Array, ' kb']
408 """The indices of binary outcome components in the full list of outcome
409 components. `None` when there are no binary components."""
411 resid: Float[Array, '*chains n'] | Float[Array, '*chains k n'] = field(
412 chains=CHAIN_AXIS, data=-1
413 )
414 """The residuals, ``resid_unit * resid = (y or z) - sum of trees``."""
416 resid_unit: Float32[Array, ''] | Float32[Array, ' k']
417 """The storage unit of `resid` (see `init`'s ``resid_dtype``), same scheme
418 as `Forest.leaf_unit`. Equal to ``leaf_unit * sqrt(num_trees)`` (the
419 marginal prior standard deviation of the sum of trees) rounded to a power
420 of two."""
422 resid_eff_scale: Float32[Array, '*chains'] | Float32[Array, '*chains k'] = field(
423 chains=CHAIN_AXIS
424 )
425 """The measured scale of the residuals (their precision-weighted root mean
426 square in data units), rounded to a power of two. Sets the leaf
427 quantization grid (see `StepConfig.leaf_quantization`). Initialized to
428 `resid_unit`, then tracks the MCMC (while the storage unit of `resid` stays
429 fixed at `resid_unit`). Not updated in purely binary regression, where the
430 scale of the latent residuals should be stable."""
432 resid_inexact_integral: Float32[Array, '*chains'] | Float32[Array, '*chains k'] = (
433 field(chains=CHAIN_AXIS)
434 )
435 """Sum over the MCMC steps done of the mean square of the residuals (in
436 `resid_unit` units) large enough that their running updates round. Used by
437 `sum_trees_eps` to estimate the accumulated rounding drift."""
439 error_cov_inv: Wishart
440 """The inverse error covariance (``error_cov_inv.value``, scalar for
441 univariate) with its Wishart prior. Fixed at the identity with no prior in
442 binary regression."""
444 error_scale: Float32[Array, ' n'] | Float32[Array, 'k n'] | None = field(data=-1)
445 """The per-datapoint error scales (the ``error_scale`` argument of `init`).
446 The error precision on a datapoint is ``error_cov_inv.value /
447 outer(error_scale, error_scale)``. For binary components the (fixed, unit)
448 probit latent error is scaled instead, so the success probability is
449 ``Phi(sum of trees / error_scale)``. `inv_sdev_scale` and `prec_scale` are
450 derived from this and the missingness mask."""
452 prec_scale: Float[Array, ' n'] | Float[Array, 'k k n'] | None = field(data=-1)
453 """The scale on the error precision, ``prec_scale = outer(inv_sdev_scale,
454 inv_sdev_scale)`` per datapoint (``inv_sdev_scale ** 2`` for scalar
455 scales), so it's in units of ``inv_sdev_unit ** 2``. Stored, like
456 `inv_sdev_scale`, in `init`'s ``prec_scale_dtype``."""
458 inv_sdev_scale: Float[Array, ' n'] | Float[Array, 'k n'] | None = field(data=-1)
459 """``inv_sdev_scale * inv_sdev_unit = 1 / error_scale``, zeroed at missing
460 datapoints. Not `None` when ``missing`` is set even if fit without error
461 scales."""
463 inv_sdev_unit: Float32[Array, ''] | Float32[Array, ' k']
464 """The storage unit of `inv_sdev_scale`, to avoid under/overflow with short
465 dtypes; same scheme as `Forest.leaf_unit`. Set to the root mean square of
466 ``1 / error_scale`` over non-missing datapoints, rounded to a power of two;
467 1 if fit without error scales. Constant along the MCMC."""
469 n_non_missing: Int32[Array, ''] | Int32[Array, ' k']
470 """The number of non-missing datapoints."""
472 sum_diag_prec_scale: Float32[Array, ''] | Float32[Array, ' k']
473 """``sum(1 / error_scale ** 2)`` over non-missing datapoints; equal to
474 `n_non_missing` if fit without error scales."""
476 forest: Forest
477 """The sum of trees model."""
479 config: StepConfig
480 """Metadata and configurations for the MCMC step."""
482 @property
483 def has_chains(self) -> bool:
484 """Whether non-constant attributes in the state carry an explicit chain axis."""
485 return self.forest.has_chains
487 def num_chains(self) -> int | None:
488 """Return the number of chains, or `None` if not multichain."""
489 if not self.has_chains:
490 return None
491 c = chain_vmap_axes(self.forest).var_tree
492 return self.forest.var_tree.shape[c]
494 @jit
495 def sum_trees_eps(self) -> Float32[Array, ''] | Float32[Array, ' k']:
496 """Estimate the absolute accuracy limit of the sum of trees (in data units).
498 The analogue of ``finfo(dtype).eps`` for the sum of trees. This combines three
499 terms: floating point resolution, random walk accumulation of numerical error
500 on running residuals, and if leaf quantization is active, breakdown of the
501 mcmc due to the quantization being too coarse. The latter term is currently
502 broken actually, sorry about that.
503 """
504 resolution, drift, snap = self._sum_trees_eps()
505 return jnp.maximum(jnp.maximum(resolution, drift), snap)
507 def _sum_trees_eps(
508 self,
509 ) -> tuple[
510 Float32[Array, ''] | Float32[Array, ' k'],
511 Float32[Array, ''] | Float32[Array, ' k'],
512 Float32[Array, ''] | Float32[Array, ' k'],
513 ]:
514 """Return the resolution, drift, and snap terms of `sum_trees_eps` separately."""
515 eps_leaf = jnp.finfo(self.forest.leaf_tree.dtype).eps
516 eps_resid = jnp.finfo(self.resid.dtype).eps
518 # rounding of the stored leaves at their typical magnitude `leaf_unit`
519 dtype_quantum = eps_leaf * self.forest.leaf_unit
521 if self.config.leaf_quantization is None:
522 resolution = dtype_quantum
523 # float leaf storage rounds relative to the leaf magnitude, so
524 # leaves near zero keep a fine grid and sampling is not pinned
525 snap = jnp.zeros_like(dtype_quantum)
526 else:
527 # quantized leaves all sit on one grid closed under addition (the
528 # scales are powers of two, so leaf storage rounds within the
529 # grid even where its spacing is coarser), so the sum resolves
530 # multiples of the quantum whatever the number of trees
531 q = self.config.leaf_quantization
532 scale_chain_axes = range(self.resid_eff_scale.ndim - self.resid_unit.ndim)
533 eff_scale = jnp.mean(self.resid_eff_scale, axis=tuple(scale_chain_axes))
534 managed_quantum = eps_resid * eff_scale * 2.0**q
535 resolution = jnp.maximum(managed_quantum, dtype_quantum)
537 # determine average error precision; the folded `inv_sdev_unit`
538 # puts the stored `prec_scale` mean in data units
539 prec = scaled_error_cov_inv(self)
540 n_eff = jnp.maximum(self.n_non_missing, 1)
541 if self.prec_scale is not None:
542 prec *= self.prec_scale.sum(axis=-1, dtype=jnp.float32) / n_eff
544 # convert precision to variance and average it over chains
545 if self.resid_unit.ndim:
546 error_var = jnp.diagonal(
547 inv_via_chol_with_gersh(prec), axis1=-2, axis2=-1
548 )
549 else:
550 error_var = jnp.reciprocal(prec)
551 var_chain_axes = range(error_var.ndim - self.resid_unit.ndim)
552 error_var = jnp.mean(error_var, axis=tuple(var_chain_axes))
554 # a quantized leaf moves only when its full conditional mean
555 # crosses half a quantum; the mean responds to the leaf's average
556 # residual through the posterior shrinkage factor s = t / (1 + t),
557 # t = (leaf_unit / error_sdev)^2 * n_leaf, so residual features
558 # smaller than quantum / (2 s) cannot move the sampler. Typical
559 # values are used for the error variance (chain mean, datapoint
560 # mean of prec_scale) and n_leaf (datapoints over leaves per tree).
561 leaves_per_tree = 1.0 + jnp.count_nonzero(self.forest.split_tree, axis=-1)
562 n_leaf = n_eff / jnp.mean(leaves_per_tree)
563 t = jnp.square(self.forest.leaf_unit) * n_leaf / error_var
564 snap = managed_quantum / 2.0 * (1.0 + jnp.reciprocal(t))
566 # random walk drift of the accumulated rounding errors: each of the
567 # `num_trees` updates per step charges eps_resid * |resid| to the
568 # residuals large enough to round, whose mean square is integrated
569 # over the steps done in `resid_inexact_integral`
570 integral = self.resid_inexact_integral
571 chain_axes = range(integral.ndim - self.resid_unit.ndim)
572 integral = jnp.mean(integral, axis=tuple(chain_axes))
573 *_, num_trees, _ = self.forest.var_tree.shape
574 drift = eps_resid * self.resid_unit * jnp.sqrt(num_trees * integral)
576 return resolution, drift, snap
579def check_diagonal(rate: Float32[Array, 'k k']) -> Float32[Array, 'k k']:
580 """Raise if the Wishart `rate` is not diagonal."""
581 diag = jnp.diag(jnp.diag(rate))
582 return error_if(rate, jnp.any(rate != diag), 'error_cov_inv.rate must be diagonal')
585def check_binary_unit_precision(
586 value: Float32[Array, 'k k'], binary_mask: Sequence[bool]
587) -> Float32[Array, 'k k']:
588 """Raise if the binary diagonal entries of `value` are not fixed at 1."""
589 binary = jnp.array(binary_mask)
590 off_unit = jnp.any(binary & (jnp.diag(value) != 1.0))
591 return error_if(
592 value,
593 off_unit,
594 'binary error precision must be 1 (the default for a zero rate)',
595 )
598def init_shape_shifting_parameters(
599 y: Float32[Array, ' n'] | Float32[Array, 'k n'],
600 outcome_type: OutcomeType | list[OutcomeType],
601 offset: Float32[Array, ''] | Float32[Array, ' k'],
602 error_scale: Float32[ArrayLike, ' n'] | Float32[ArrayLike, 'k n'] | None,
603 error_cov_inv: Wishart | None,
604 leaf_prior_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'],
605 missing: Bool[ArrayLike, ' n'] | Bool[ArrayLike, 'k n'] | None,
606) -> tuple[bool, tuple[int, ...], Wishart, None | Int32[Array, ' kb']]:
607 """
608 Check and initialize parameters that change array type/shape based on outcome kind.
610 Parameters
611 ----------
612 y
613 The response variable (used only for shape checks).
614 outcome_type
615 Whether the regression is continuous or binary. Can be a list of
616 `OutcomeType` for per-component specification in the multivariate case.
617 offset
618 The offset to add to the predictions.
619 error_scale
620 Per-observation error scale (univariate only).
621 error_cov_inv
622 The Wishart prior on the error precision and its initial value, or
623 `None` for binary regression. The mixed and partial-missing diagonal
624 modes require a `DiagWishart`; in the mixed case the binary components
625 must have an initial precision of 1.
626 leaf_prior_cov_inv
627 The inverse of the leaf prior covariance.
628 missing
629 The per-datapoint missingness mask, used to detect partial missingness
630 (2-D mask) so that diagonal-mode initialization is selected.
632 Returns
633 -------
634 is_binary
635 Whether all outcomes are binary.
636 kshape
637 The outcome shape, empty for univariate, (k,) for multivariate.
638 error_cov_inv
639 The Wishart prior with its initial value resolved for the outcome kind.
640 binary_indices
641 The indices of binary outcome components, or `None` if there are none.
642 """
643 kshape = offset.shape
645 # determine per-component outcome kinds
646 if isinstance(outcome_type, list):
647 assert kshape, 'per-component outcome_type requires multivariate y'
648 (k,) = kshape
649 assert len(outcome_type) == k
650 binary_mask = [t is OutcomeType.binary for t in outcome_type]
651 is_binary = all(binary_mask)
652 is_mixed = any(binary_mask) and not is_binary
653 else:
654 is_binary = outcome_type is OutcomeType.binary
655 is_mixed = False
657 if is_mixed:
658 binary_indices = jnp.array([i for i, b in enumerate(binary_mask) if b])
659 else:
660 binary_indices = None
662 partial_missing = missing is not None and missing.ndim == 2 and kshape
664 assert (
665 error_scale is None
666 or error_scale.shape == y.shape # (k, n)
667 or error_scale.shape == y.shape[-1:] # (n,)
668 )
670 # All-binary: no prior, the precision is fixed at the identity.
671 if is_binary:
672 assert error_cov_inv is None, 'no error covariance prior in binary regression'
673 value = jnp.eye(kshape[0]) if kshape else jnp.array(1.0)
674 error_cov_inv = Wishart(nu=None, rate=None, value=value)
676 # Mixed binary-continuous, or continuous-mv with 2-D missingness: diagonal
677 # error covariance, updated component-wise. The caller must supply a
678 # `DiagWishart`; in the mixed case the binary components must have unit
679 # initial precision (see `DiagWishart`).
680 elif is_mixed or partial_missing:
681 assert isinstance(error_cov_inv, DiagWishart), (
682 'mixed binary-continuous or partial-missing regression requires a '
683 'DiagWishart error_cov_inv prior'
684 )
685 assert error_cov_inv.rate is not None
686 assert error_cov_inv.rate.shape == 2 * kshape
687 assert error_cov_inv.value.shape == 2 * kshape
688 rate = check_diagonal(error_cov_inv.rate)
689 value = check_diagonal(error_cov_inv.value)
690 if is_mixed:
691 value = check_binary_unit_precision(value, binary_mask)
692 error_cov_inv = replace(error_cov_inv, rate=rate, value=value)
694 # All-continuous: a dense `Wishart`.
695 else:
696 assert error_cov_inv is not None
697 assert type(error_cov_inv) is Wishart, (
698 'continuous regression requires a dense Wishart error_cov_inv prior'
699 )
700 rate = error_cov_inv.rate
701 assert rate is not None
702 assert rate.shape == 2 * kshape
703 assert error_cov_inv.value.shape == 2 * kshape
705 assert y.shape[:-1] == kshape
706 assert leaf_prior_cov_inv.shape == 2 * kshape
708 return is_binary, kshape, error_cov_inv, binary_indices
711def check_splitless_vars(
712 filter_splitless_vars: int,
713 max_split: UInt[Array, ' p'],
714 offset: Float32[Array, ''] | Float32[Array, ' k'],
715) -> Float32[Array, ''] | Float32[Array, ' k']:
716 """Check there aren't too many deactivated predictors."""
717 msg = (
718 f'there are more than {filter_splitless_vars=} predictors with no splits, '
719 'please increase `filter_splitless_vars` or investigate the missing splits'
720 )
721 return error_if(offset, jnp.sum(max_split == 0) > filter_splitless_vars, msg)
724def parse_outcome_type(
725 outcome_type: 'OutcomeType | str | Sequence[OutcomeType | str]',
726) -> 'OutcomeType | list[OutcomeType]':
727 """Normalize outcome_type to enum (or list of enums)."""
728 if isinstance(outcome_type, Sequence) and not isinstance(outcome_type, str):
729 return [OutcomeType(t) for t in outcome_type]
730 else:
731 return OutcomeType(outcome_type)
734def parse_p_nonterminal(
735 p_nonterminal: Float32[ArrayLike, ' d_minus_1'],
736) -> Float32[Array, ' d_minus_1+1']:
737 """Check it's in (0, 1) and pad with a 0 at the end."""
738 p_nonterminal = jnp.asarray(p_nonterminal)
739 ok = (p_nonterminal > 0) & (p_nonterminal < 1)
740 p_nonterminal = error_if(p_nonterminal, ~ok, 'p_nonterminal must be in (0, 1)')
741 return jnp.pad(p_nonterminal, (0, 1))
744def make_p_nonterminal(
745 d: int, alpha: FloatLike = 0.95, beta: FloatLike = 2.0
746) -> Float32[Array, ' {d}-1']:
747 """Prepare the `p_nonterminal` argument to `init`.
749 It is calculated according to the formula:
751 P_nt(depth) = alpha / (1 + depth)^beta, with depth 0-based
753 Parameters
754 ----------
755 d
756 The maximum depth of the trees (d=1 means tree with only root node)
757 alpha
758 The a priori probability of the root node having children, conditional
759 on it being possible
760 beta
761 The exponent of the power decay of the probability of having children
762 with depth.
764 Returns
765 -------
766 An array of probabilities, one per tree level but the last.
767 """
768 assert d >= 1
769 depth = jnp.arange(d - 1)
770 return alpha / (1 + depth).astype(float) ** beta
773@float32_matmuls
774def init(
775 *,
776 X: UInt[ArrayLike, 'p n'],
777 y: Float32[ArrayLike, ' n'] | Float32[ArrayLike, ' k n'],
778 outcome_type: OutcomeType | str | Sequence[OutcomeType | str] = 'continuous',
779 offset: FloatLike | Float[ArrayLike, ' k'],
780 max_split: UInt[ArrayLike, ' p'],
781 num_trees: int,
782 p_nonterminal: Float32[ArrayLike, ' d_minus_1'],
783 leaf_prior_cov_inv: FloatLike | Float[ArrayLike, 'k k'],
784 leaf_dtype: DTypeLike = jnp.float32,
785 prec_scale_dtype: DTypeLike = jnp.float32,
786 resid_dtype: DTypeLike = jnp.float32,
787 leaf_quantization: int | Integer[ArrayLike, ''] | None = None,
788 error_cov_inv: Wishart | None = None,
789 error_scale: Float32[ArrayLike, ' n'] | Float32[ArrayLike, 'k n'] | None = None,
790 missing: Bool[ArrayLike, ' n'] | Bool[ArrayLike, 'k n'] | None = None,
791 min_points_per_decision_node: int | Integer[ArrayLike, ''] | None = None,
792 min_points_per_leaf: int | Integer[ArrayLike, ''] | None = None,
793 resid_reduction_config: ReductionConfig = AutoBatchedReduction(),
794 count_reduction_config: ReductionConfig = AutoOneHotReduction(),
795 prec_reduction_config: ReductionConfig = AutoOneHotReduction(),
796 prec_count_num_trees: int | None | Literal['auto'] = 'auto',
797 sequential_unroll: int | bool = 2,
798 save_ratios: bool = False,
799 filter_splitless_vars: int = 0,
800 log_s: Float32[ArrayLike, ' p'] | None = None,
801 theta: FloatLike | None = None,
802 a: FloatLike | None = None,
803 b: FloatLike | None = None,
804 rho: FloatLike | None = None,
805 sparse_on_at: int | Integer[ArrayLike, ''] | None = None,
806 augment: bool = True,
807 num_chains: int | None = None,
808 mesh: Mesh | dict[str, int] | None = None,
809) -> State:
810 """
811 Make a BART posterior sampling MCMC initial state.
813 Parameters
814 ----------
815 X
816 The predictors. Note this is trasposed compared to the usual convention.
817 y
818 The response. If two-dimensional, the outcome is multivariate with the
819 first axis indicating the component. For binary data, non-zero means 1,
820 zero means 0.
821 outcome_type
822 Whether the regression is continuous or binary (probit). Can also be a
823 sequence of `OutcomeType` values, one per outcome component, for mixed
824 binary-continuous multivariate regression.
825 offset
826 Constant shift added to the sum of trees. 0 if not specified.
827 max_split
828 The maximum split index for each variable. All split ranges start at 1.
829 num_trees
830 The number of trees in the forest.
831 p_nonterminal
832 The probability of a nonterminal node at each depth. The maximum depth
833 of trees is fixed by the length of this array. Use `make_p_nonterminal`
834 to set it with the conventional formula.
835 leaf_prior_cov_inv
836 The prior precision matrix of a leaf, see `Forest.leaf_prior_cov_inv`.
837 leaf_dtype
838 prec_scale_dtype
839 resid_dtype
840 Storage dtypes for, respectively: leaves (`Forest.leaf_tree`), derived
841 error scales (`State.prec_scale` and `State.inv_sdev_scale`; the raw
842 `State.error_scale` stays float32), and running residuals
843 (`State.resid`). These quantities are stored in scaled units so that
844 narrow dtypes do not under/overflow (see `Forest.leaf_unit`), though
845 float16 error scales may still overflow if their dynamic range is high
846 enough. Leaf full conditionals are always computed and sampled in
847 float32. Narrow residual dtypes may easily break the MCMC,
848 `resid_dtype` is an experimental setting.
849 leaf_quantization
850 Quantize the leaves to (almost) stop the numerical drift of the running
851 residuals, see `StepConfig.leaf_quantization`.
852 error_cov_inv
853 The prior and initial value of the error term precision matrix (see
854 `Wishart`). Leave it unspecified for binary regression. Mixed
855 binary-continuous and per-outcome-component missingness require a
856 `DiagWishart`.
857 error_scale
858 Per-datapoint error scales (called ``w`` in the R package BART3); see
859 `State.error_scale`. An unspecified `error_scale` is equivalent to
860 ``error_scale = 1`` for all datapoints.
861 missing
862 Boolean mask indicating which datapoints are missing. `True` marks
863 entries to be ignored by the MCMC. The masked values of `y` may be
864 anything, even non-finite. If `missing` is 2-D, `error_cov_inv` must
865 be a `DiagWishart`.
866 min_points_per_decision_node
867 min_points_per_leaf
868 The minimum number of datapoints in a decision node and in a leaf,
869 respectively; 0 if not specified. The leaf constraint is not taken into
870 account in the proposal distribution because that would be expensive.
871 The two are independent and not checked for coherence; it makes sense
872 to set ``min_points_per_decision_node >= 2 * min_points_per_leaf``.
873 resid_reduction_config
874 count_reduction_config
875 prec_reduction_config
876 How to sum the residuals, count the datapoints, and sum the likelihood
877 precisions in each leaf, respectively. See `ReductionConfig` and its
878 subclasses.
879 prec_count_num_trees
880 The number of trees to process at a time when counting datapoints or
881 computing the likelihood precision. If `None`, do all trees at once,
882 which may use too much memory on cpu. If 'auto' (default), it's chosen
883 automatically.
884 sequential_unroll
885 See `StepConfig.sequential_unroll`. Unrolling may speed up the MCMC at
886 the cost of longer compilation; 1 means no unrolling.
887 save_ratios
888 Whether to save the Metropolis-Hastings ratios.
889 filter_splitless_vars
890 The maximum number of variables without splits that can be ignored. If
891 there are more, `init` raises an exception.
892 log_s
893 theta
894 a
895 b
896 rho
897 Sparsity (variable selection) parameters, see `Forest.log_s` and
898 `Forest.theta`. If `rho`, `a`, `b` are set, an unspecified `theta` is
899 initialized to `rho`; if `theta` is set, an unspecified `log_s` is
900 initialized to uniform.
901 sparse_on_at
902 After how many MCMC steps to turn on variable selection.
903 augment
904 See `StepConfig.augment`. If disabled, forbidden decision rules are
905 ignored when counting variable usage, which may be faster but is
906 an approximation.
907 num_chains
908 The number of independent MCMC chains. Single chain with scalar values
909 if not specified.
910 mesh
911 A jax mesh used to shard data and computation across multiple devices.
912 If it has a 'chains' axis, that axis is used to shard the chains. If it
913 has a 'data' axis, that axis is used to shard the datapoints.
915 As a shorthand, if a dictionary mapping axis names to axis size is
916 passed, the corresponding mesh is created, e.g., ``dict(chains=4,
917 data=2)`` will let jax pick 8 devices to split chains (which must be a
918 multiple of 4) across 4 pairs of devices, where in each pair the data
919 is split in two.
921 Note: if a mesh is passed, the arrays are always sharded according to
922 it. In particular even if the mesh has no 'chains' or 'data' axis, the
923 arrays will be replicated on all devices in the mesh.
925 Returns
926 -------
927 An initialized BART MCMC state.
929 Raises
930 ------
931 ValueError
932 If arguments unused in binary regression are set.
934 Notes
935 -----
936 In decision nodes, the values in ``X[i, :]`` are compared to a cutpoint out
937 of the range ``[1, 2, ..., max_split[i]]``. A point belongs to the left
938 child iff ``X[i, j] < cutpoint``. Thus it makes sense for ``X[i, :]`` to be
939 integers in the range ``[0, 1, ..., max_split[i]]``.
941 In general the arrays passed to this function as arguments may be donated,
942 invalidating them. Create copies before passing them to `init` if this
943 happens and you need them again.
944 """
945 # convert to array all array-like arguments that are used in other
946 # configurations but don't need further processing themselves
947 X = jnp.asarray(X)
948 y = jnp.asarray(y)
949 assert y.dtype == jnp.float32
950 offset = jnp.asarray(offset)
951 leaf_prior_cov_inv = jnp.asarray(leaf_prior_cov_inv)
952 max_split = jnp.asarray(max_split)
953 error_scale = asarray_or_none(error_scale)
954 missing = asarray_or_none(missing)
955 assert missing is None or missing.ndim <= y.ndim
957 # normalize outcome_type to enum (or list of enums)
958 outcome_type = parse_outcome_type(outcome_type)
960 # check p_nonterminal and pad it with a 0 at the end (still not final shape)
961 p_nonterminal = parse_p_nonterminal(p_nonterminal)
963 # process arguments that change depending on outcome type
964 is_binary, kshape, error_cov_inv, binary_indices = init_shape_shifting_parameters(
965 y, outcome_type, offset, error_scale, error_cov_inv, leaf_prior_cov_inv, missing
966 )
968 storage = determine_storage_params(
969 leaf_dtype, prec_scale_dtype, resid_dtype, leaf_prior_cov_inv, kshape, num_trees
970 )
972 # extract array sizes from arguments
973 (max_depth,) = p_nonterminal.shape
974 p, n = X.shape
976 # check and initialize sparsity parameters
977 if not all_none_or_not_none(rho, a, b): 977 ↛ 978line 977 didn't jump to line 978 because the condition on line 977 was never true
978 msg = 'rho, a, b are not either all `None` or all set'
979 raise ValueError(msg)
980 if theta is None and rho is not None:
981 theta = rho
982 if log_s is None and theta is not None:
983 log_s = jnp.zeros(max_split.size)
984 if not all_none_or_not_none(theta, sparse_on_at): 984 ↛ 985line 984 didn't jump to line 985 because the condition on line 984 was never true
985 msg = 'sparsity params (either theta or rho,a,b) and sparse_on_at must be either all None or all set'
986 raise ValueError(msg)
988 # determine settings for reductions
989 mesh = parse_mesh(num_chains, mesh)
990 red_cfg = parse_reduction_configs(
991 resid_reduction_config,
992 count_reduction_config,
993 prec_reduction_config,
994 prec_count_num_trees,
995 y,
996 num_trees,
997 num_chains,
998 mesh,
999 )
1001 # check there aren't too many deactivated predictors
1002 offset = check_splitless_vars(filter_splitless_vars, max_split, offset)
1004 tree_size = 2**max_depth
1006 # Assemble the state, shard it, then fill in the post-shard fields. This
1007 # whole region runs with type-checking disabled because the state carries
1008 # deliberately wrong-typed intermediates parked in its fields for sharding:
1009 # `_LazyArray` leaves (each chain-bearing leaf is built at its core no-chain
1010 # shape, then `add_chains` wraps it to broadcast in the chain axis) and the
1011 # user `missing` mask in the `inv_sdev_scale` slot. The context ends once
1012 # every field has been replaced by its final, correctly-typed array.
1013 with jaxtyping_disabled():
1014 state = State(
1015 _chain_anchor=lazy(jnp.zeros, ()), # typechecker chain anchor
1016 X=X,
1017 y=y,
1018 z=(
1019 lazy(jnp.full, y.shape, offset[..., None])
1020 if is_binary
1021 else lazy(
1022 jnp.full, (binary_indices.size, n), offset[binary_indices, None]
1023 )
1024 if binary_indices is not None
1025 else None
1026 ),
1027 binary_indices=binary_indices,
1028 resid=(
1029 lazy(jnp.zeros, y.shape, storage.resid_dtype)
1030 if is_binary
1031 # resid is created later after y and offset are sharded
1032 else cast(Array, None)
1033 ),
1034 resid_unit=storage.resid_unit,
1035 resid_eff_scale=lazy(jnp.full, kshape, storage.resid_unit),
1036 resid_inexact_integral=lazy(jnp.zeros, kshape),
1037 # only `value` carries the chain axis, so it becomes the lazy leaf;
1038 # the prior params `nu`/`rate` are shared across chains
1039 error_cov_inv=replace(
1040 error_cov_inv, value=_lazy_from_array(error_cov_inv.value)
1041 ),
1042 # `error_scale` goes straight to its field; `missing` is parked in
1043 # the `inv_sdev_scale` slot so it gets sharded with everything
1044 # else. `compute_scale_related_attrs` derives `prec_scale` and
1045 # `inv_sdev_scale` post-shard.
1046 error_scale=error_scale,
1047 prec_scale=None,
1048 inv_sdev_scale=missing,
1049 # invalid placeholders; the true values are set post-shard by
1050 # `compute_scale_related_attrs`
1051 inv_sdev_unit=cast(Array, None),
1052 n_non_missing=cast(Array, None),
1053 sum_diag_prec_scale=cast(Array, None),
1054 forest=Forest(
1055 leaf_tree=lazy(
1056 jnp.zeros, (num_trees, *kshape, tree_size), storage.leaf_dtype
1057 ),
1058 leaf_unit=storage.leaf_unit,
1059 offset=offset,
1060 var_tree=lazy(
1061 jnp.zeros,
1062 (num_trees, tree_size // 2),
1063 minimal_unsigned_dtype(p - 1),
1064 ),
1065 split_tree=lazy(
1066 jnp.zeros, (num_trees, tree_size // 2), max_split.dtype
1067 ),
1068 affluence_tree=lazy(
1069 initial_affluence_tree,
1070 (num_trees, tree_size // 2),
1071 n,
1072 min_points_per_decision_node,
1073 ),
1074 blocked_vars=get_blocked_vars(filter_splitless_vars, max_split),
1075 max_split=max_split,
1076 grow_prop_count=lazy(jnp.zeros, (), int),
1077 grow_acc_count=lazy(jnp.zeros, (), int),
1078 prune_prop_count=lazy(jnp.zeros, (), int),
1079 prune_acc_count=lazy(jnp.zeros, (), int),
1080 p_nonterminal=p_nonterminal[tree_depths(tree_size)],
1081 p_propose_grow=p_nonterminal[tree_depths(tree_size // 2)],
1082 leaf_indices=lazy(
1083 jnp.ones, (num_trees, n), minimal_unsigned_dtype(tree_size - 1)
1084 ),
1085 # no pending prune: `step` starts by applying the pending
1086 # prunes to `leaf_indices`, which shall be a no-op on init
1087 to_prune=lazy(jnp.zeros, (num_trees,), bool),
1088 move_node=lazy(jnp.zeros, (num_trees,), jnp.int32),
1089 # the counts serve the minimum-points constraints and stand in
1090 # for the precisions when there are no per-datapoint scales
1091 # (`prec_scale` is set iff `error_scale` or `missing` is given)
1092 count_tree=(
1093 lazy(initial_count_tree, (num_trees, tree_size), n)
1094 if min_points_per_decision_node is not None
1095 or min_points_per_leaf is not None
1096 or (error_scale is None and missing is None)
1097 else None
1098 ),
1099 # prec_tree is created later, it needs the sharded prec_scale
1100 prec_tree=None,
1101 min_points_per_decision_node=asarray_or_none(
1102 min_points_per_decision_node
1103 ),
1104 min_points_per_leaf=asarray_or_none(min_points_per_leaf),
1105 log_trans_prior=lazy(jnp.zeros, (num_trees,)) if save_ratios else None,
1106 log_likelihood=lazy(jnp.zeros, (num_trees,)) if save_ratios else None,
1107 leaf_prior_cov_inv=leaf_prior_cov_inv,
1108 log_s=_lazy_from_array(asarray_or_none(log_s)),
1109 theta=_lazy_from_array(asarray_or_none(theta)),
1110 rho=asarray_or_none(rho),
1111 a=asarray_or_none(a),
1112 b=asarray_or_none(b),
1113 ),
1114 config=StepConfig(
1115 steps_done=jnp.int32(0),
1116 sparse_on_at=asarray_or_none(sparse_on_at),
1117 sequential_unroll=sequential_unroll,
1118 augment=augment,
1119 mesh=mesh,
1120 leaf_quantization=asarray_or_none(leaf_quantization),
1121 **red_cfg,
1122 ),
1123 )
1125 # add the chain axis to every chain-marked leaf at the position
1126 # declared by its field metadata
1127 state = add_chains(state, num_chains)
1129 # delete big input arrays such that they can be deleted as soon as they
1130 # are sharded, only those arrays that contain an (n,) sized axis
1131 del X, error_scale, missing, y
1133 # move all arrays to the appropriate device and instantiate lazy arrays
1134 state = shard_state(state)
1136 # replace y at masked positions post-shard (the mask is parked in
1137 # `inv_sdev_scale`), before `resid` is derived from y
1138 if state.inv_sdev_scale is not None:
1139 state = replace(
1140 state, y=sanitize_y(state.y, state.inv_sdev_scale, state.forest.offset)
1141 )
1143 # derive the scale arrays and their constant summaries after sharding
1144 # to do the calculation on the right devices. `state.error_scale`
1145 # already holds the sharded user-supplied scale and
1146 # `state.inv_sdev_scale` holds the parked `missing` mask.
1147 # `compute_scale_related_attrs` does not donate `error_scale`, so it
1148 # stays in place; the derived scales fold in the mask, the raw scale
1149 # does not.
1150 attrs = compute_scale_related_attrs(
1151 state.error_scale, state.inv_sdev_scale, storage.prec_scale_dtype, n
1152 )
1153 state = replace(
1154 state,
1155 prec_scale=attrs.prec_scale,
1156 inv_sdev_scale=attrs.inv_sdev_scale,
1157 inv_sdev_unit=attrs.inv_sdev_unit,
1158 n_non_missing=attrs.n_non_missing,
1159 sum_diag_prec_scale=attrs.sum_diag_prec_scale,
1160 )
1162 # calculate initial resid in the continuous outcome case, such that y
1163 # and offset are already sharded if needed
1164 if state.resid is None:
1165 state = set_initial_resid(
1166 state, binary_indices, num_chains, storage.resid_dtype
1167 )
1168 # charge the one-time rounding of the initial residuals into storage
1169 # when it uses a dtype narrower than y (else the cast is exact)
1170 if jnp.finfo(storage.resid_dtype).nmant < jnp.finfo(state.y.dtype).nmant:
1171 state = replace(
1172 state,
1173 resid_inexact_integral=initial_resid_inexact_integral(
1174 state.resid, state.n_non_missing, num_trees
1175 ),
1176 )
1178 # calculate the initial prec_tree from the sharded prec_scale
1179 if state.prec_scale is not None:
1180 state = set_initial_prec_tree(state, num_chains, num_trees, tree_size)
1182 # all the wrong-typed intermediates have now been replaced by their final
1183 # values, so type-checking can resume; make all types strong to avoid
1184 # unwanted recompilations
1185 return remove_weak_types(state)
1188def parse_float_dtype(dtype: DTypeLike) -> jnp.dtype:
1189 """Normalize a storage dtype and check it is floating point."""
1190 dtype = jnp.dtype(dtype)
1191 assert jnp.issubdtype(dtype, jnp.floating)
1192 return dtype
1195def compute_leaf_unit(
1196 leaf_prior_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'],
1197 kshape: tuple[int, ...],
1198) -> Float32[Array, ''] | Float32[Array, ' k']:
1199 """Compute the marginal prior standard deviation of a leaf.
1201 A degenerate prior precision (e.g., infinite, from `bartz.Bart` with
1202 constant ``y``) yields a zero or non-finite scale; fall back to 1 to avoid
1203 nan leaves (an infinite precision pins the leaves to zero anyway).
1204 """
1205 if kshape:
1206 leaf_prior_cov = inv_via_chol_with_gersh(leaf_prior_cov_inv)
1207 leaf_unit = jnp.sqrt(jnp.diagonal(leaf_prior_cov))
1208 else:
1209 leaf_unit = jnp.sqrt(jnp.reciprocal(leaf_prior_cov_inv))
1210 return jnp.where(jnp.isfinite(leaf_unit) & (leaf_unit > 0), leaf_unit, 1.0)
1213def round_to_pow2(
1214 x: Float32[Array, ''] | Float32[Array, ' k'],
1215) -> Float32[Array, ''] | Float32[Array, ' k']:
1216 """Round to the nearest power of two."""
1217 # note: don't use exp2, not exact. `2 ** x` checked exact on cpu & cuda.
1218 return 2 ** jnp.round(jnp.log2(x))
1221def scaled_error_cov_inv(
1222 state: State,
1223) -> Float32[Array, '*chains'] | Float32[Array, '*chains k k']:
1224 """Return the error precision with ``inv_sdev_unit ** 2`` folded in.
1226 `State.prec_scale` and everything summed from it (`Forest.prec_tree`, the
1227 per-leaf precision-scaled residual sums) are stored in units of
1228 ``State.inv_sdev_unit ** 2``, so their products with this scaled precision
1229 are in data units. Returns ``error_cov_inv.value`` as-is when there are no
1230 per-datapoint error scales.
1231 """
1232 value = state.error_cov_inv.value
1233 unit = state.inv_sdev_unit
1234 if state.prec_scale is None:
1235 return value
1236 elif unit.ndim:
1237 return value * (unit[:, None] * unit[None, :])
1238 else:
1239 return value * jnp.square(unit)
1242@dataclass(frozen=True)
1243class StorageParams:
1244 """Storage dtypes and units for the leaves and residuals."""
1246 leaf_dtype: jnp.dtype
1247 prec_scale_dtype: jnp.dtype
1248 resid_dtype: jnp.dtype
1249 leaf_unit: Float32[Array, ''] | Float32[Array, ' k']
1250 resid_unit: Float32[Array, ''] | Float32[Array, ' k']
1253def determine_storage_params(
1254 leaf_dtype: DTypeLike,
1255 prec_scale_dtype: DTypeLike,
1256 resid_dtype: DTypeLike,
1257 leaf_prior_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'],
1258 kshape: tuple[int, ...],
1259 num_trees: int,
1260) -> StorageParams:
1261 """Normalize the storage dtypes and compute the leaf and residual units."""
1262 leaf_unit = compute_leaf_unit(leaf_prior_cov_inv, kshape)
1263 return StorageParams(
1264 leaf_dtype=parse_float_dtype(leaf_dtype),
1265 prec_scale_dtype=parse_float_dtype(prec_scale_dtype),
1266 resid_dtype=parse_float_dtype(resid_dtype),
1267 leaf_unit=round_to_pow2(leaf_unit),
1268 resid_unit=round_to_pow2(leaf_unit * num_trees**0.5),
1269 )
1272def set_initial_resid(
1273 state: 'State',
1274 binary_indices: Int32[Array, ' kb'] | None,
1275 num_chains: int | None,
1276 resid_dtype: jnp.dtype,
1277) -> 'State':
1278 """Build the continuous-outcome `resid` and shard it.
1280 Called post-shard so the captured ``state.y`` and
1281 ``state.forest.offset`` are already on the target devices. Sharding axes are
1282 read via `chain_vmap_axes` / `data_vmap_axes` on a shape preview where the
1283 new `resid` leaf has the chain-extended ``ndim`` (inflated by a placeholder
1284 when `num_chains` is not `None`).
1285 """
1286 inner = _LazyArray(
1287 initial_resid,
1288 state.y.shape,
1289 state.y,
1290 state.forest.offset,
1291 binary_indices,
1292 state.resid_unit,
1293 resid_dtype,
1294 )
1295 preview_resid = add_dummy_axis(inner) if num_chains is not None else inner
1296 preview = replace(state, resid=preview_resid)
1297 chain_axis = chain_vmap_axes(preview).resid
1298 data_axis = data_vmap_axes(preview).resid
1299 resid = _wrap_chain(inner, chain_axis, num_chains)
1300 resid = shard_leaf(resid, chain_axis, data_axis, state.config.mesh)
1301 return replace(state, resid=resid)
1304@jit
1305def initial_resid_inexact_integral(
1306 resid: Float[Array, '*chains n'] | Float[Array, '*chains k n'],
1307 n_non_missing: Int32[Array, ''] | Int32[Array, ' k'],
1308 num_trees: int,
1309) -> Float32[Array, '*chains'] | Float32[Array, '*chains k']:
1310 """Seed value for `State.resid_inexact_integral` from the initial rounding.
1312 Casting the initial residuals to a storage dtype narrower than `y` rounds
1313 them once, and the running updates never fix this offset (they either round
1314 again, which the per-step accounting covers, or preserve it exactly), so it
1315 is charged upfront as one tree-update's worth of rounding.
1316 """
1317 # masked residuals are set to 0 at this stage, so they drop out of the sum
1318 ms = jnp.einsum('...n,...n->...', resid, resid, preferred_element_type=jnp.float32)
1319 ms /= jnp.maximum(n_non_missing, 1)
1320 return ms / num_trees
1323def initial_resid(
1324 shape: tuple[int, ...],
1325 y: Float32[Array, ' n'] | Float32[Array, 'k n'],
1326 offset: Float32[Array, ''] | Float32[Array, ' k'],
1327 binary_indices: Int32[Array, ' kb'] | None,
1328 resid_unit: Float32[Array, ''] | Float32[Array, ' k'],
1329 resid_dtype: jnp.dtype,
1330) -> Float[Array, ' n'] | Float[Array, 'k n']:
1331 """Calculate the initial value for `State.resid` in the continuous outcome case.
1333 The residual is stored in units of `resid_unit` and dtype `resid_dtype`. In
1334 the mixed binary-continuous case, binary rows are zeroed out (their residual
1335 starts at ``z - trees - offset = 0``).
1336 """
1337 resid = jnp.broadcast_to(y - offset[..., None], shape)
1338 if binary_indices is not None:
1339 resid = resid.at[..., binary_indices, :].set(0.0)
1340 return (resid / resid_unit[..., None]).astype(resid_dtype)
1343def initial_affluence_tree(
1344 shape: tuple[int, ...], n: int, min_points_per_decision_node: int | None
1345) -> Shaped[Array, '...']:
1346 """Create the initial value of `Forest.affluence_tree`."""
1347 return (
1348 jnp.zeros(shape, bool)
1349 .at[..., 1]
1350 .set(
1351 True
1352 if min_points_per_decision_node is None
1353 else n >= min_points_per_decision_node
1354 )
1355 )
1358def initial_count_tree(shape: tuple[int, ...], n: int) -> Shaped[Array, '...']:
1359 """Create the initial value of `Forest.count_tree`: all datapoints in the root."""
1360 return jnp.zeros(shape, jnp.uint32).at[..., 1].set(n)
1363def set_initial_prec_tree(
1364 state: State, num_chains: int | None, num_trees: int, tree_size: int
1365) -> State:
1366 """Build the cached per-leaf precision for root-only trees and shard it.
1368 Called post-shard so the captured ``state.prec_scale`` is already on the
1369 target devices; mirrors `set_initial_resid`.
1370 """
1371 assert state.prec_scale is not None
1372 shape = (num_trees, *state.prec_scale.shape[:-1], tree_size)
1373 inner = _LazyArray(initial_prec_tree, shape, state.prec_scale)
1374 preview_tree = add_dummy_axis(inner) if num_chains is not None else inner
1375 preview = replace(state, forest=replace(state.forest, prec_tree=preview_tree))
1376 chain_axis = chain_vmap_axes(preview).forest.prec_tree
1377 prec_tree = _wrap_chain(inner, chain_axis, num_chains)
1378 prec_tree = shard_leaf(prec_tree, chain_axis, None, state.config.mesh)
1379 return replace(state, forest=replace(state.forest, prec_tree=prec_tree))
1382def initial_prec_tree(
1383 shape: tuple[int, ...], prec_scale: Float[Array, ' n'] | Float[Array, 'k k n']
1384) -> Float32[Array, 'num_trees tree_size'] | Float32[Array, 'num_trees k k tree_size']:
1385 """Create the initial value of `Forest.prec_tree`: all datapoints in the root."""
1386 return (
1387 jnp.zeros(shape, jnp.float32)
1388 .at[..., 1]
1389 .set(prec_scale.sum(axis=-1, dtype=jnp.float32))
1390 )
1393@jit(donate_argnums=(0,))
1394def sanitize_y(
1395 y: Float32[Array, ' n'] | Float32[Array, 'k n'],
1396 missing: Bool[Array, ' n'] | Bool[Array, 'k n'],
1397 offset: Float32[Array, ''] | Float32[Array, ' k'],
1398) -> Float32[Array, ' n'] | Float32[Array, 'k n']:
1399 """Replace `y` with `offset` at masked positions.
1401 The MCMC ignores masked datapoints through their zeroed precision, but the
1402 values parked in `y` still enter `resid`; garbage values of large magnitude
1403 would degrade the accuracy of quantities derived from it, like the ``y -
1404 resid`` train predictions. Donates `y` to overwrite it in place.
1405 """
1406 return jnp.where(missing, offset[..., None], y)
1409class ScaleRelatedAttrs(NamedTuple):
1410 """Output of `compute_scale_related_attrs`."""
1412 inv_sdev_scale: Float[Array, ' n'] | Float[Array, 'k n'] | None
1413 prec_scale: Float[Array, ' n'] | Float[Array, 'k k n'] | None
1414 inv_sdev_unit: Float32[Array, ''] | Float32[Array, ' k']
1415 n_non_missing: Int32[Array, ''] | Int32[Array, ' k']
1416 sum_diag_prec_scale: Float32[Array, ''] | Float32[Array, ' k']
1419@jit(donate_argnums=(1,), static_argnums=(2, 3))
1420def compute_scale_related_attrs(
1421 error_scale: Float32[Array, ' n'] | Float32[Array, 'k n'] | None,
1422 missing: Bool[Array, ' n'] | Bool[Array, 'k n'] | None,
1423 prec_scale_dtype: jnp.dtype,
1424 n: int,
1425) -> ScaleRelatedAttrs:
1426 """Compute all the fixed quantities derived from `error_scale` and `missing`."""
1427 if error_scale is None and missing is None:
1428 n_non_missing = jnp.full((), n)
1429 return ScaleRelatedAttrs(
1430 inv_sdev_scale=None,
1431 prec_scale=None,
1432 inv_sdev_unit=jnp.float32(1.0),
1433 n_non_missing=n_non_missing,
1434 sum_diag_prec_scale=n_non_missing.astype(jnp.float32),
1435 )
1437 # compute inv_sdev_scale
1438 if error_scale is None: 1438 ↛ 1439line 1438 didn't jump to line 1439 because the condition on line 1438 was never true
1439 inv_sdev_scale = jnp.array(1.0) # becomes a vector with the `where` below
1440 else:
1441 inv_sdev_scale = jnp.reciprocal(error_scale)
1442 if missing is not None:
1443 inv_sdev_scale = jnp.where(missing, 0.0, inv_sdev_scale)
1445 # count non-missing points, and the equivalent error-scaled quantity
1446 n_non_missing = jnp.sum(inv_sdev_scale != 0, axis=-1)
1447 sum_diag_prec_scale = jnp.einsum('...n,...n->...', inv_sdev_scale, inv_sdev_scale)
1449 # compute inv_sdev_unit and factor it out of inv_sdev_scale
1450 inv_sdev_unit = round_to_pow2(
1451 jnp.sqrt(sum_diag_prec_scale / jnp.maximum(n_non_missing, 1))
1452 )
1453 inv_sdev_unit = jnp.where(inv_sdev_unit, inv_sdev_unit, 1.0)
1454 inv_sdev_scale = inv_sdev_scale / inv_sdev_unit[..., None]
1456 # compute prec_scale
1457 if inv_sdev_scale.ndim == 1:
1458 prec_scale = jnp.square(inv_sdev_scale)
1459 else:
1460 prec_scale = jnp.einsum('an,bn->abn', inv_sdev_scale, inv_sdev_scale)
1462 return ScaleRelatedAttrs(
1463 inv_sdev_scale=inv_sdev_scale.astype(prec_scale_dtype),
1464 prec_scale=prec_scale.astype(prec_scale_dtype),
1465 inv_sdev_unit=inv_sdev_unit,
1466 n_non_missing=n_non_missing,
1467 sum_diag_prec_scale=sum_diag_prec_scale,
1468 )
1471def get_blocked_vars(
1472 filter_splitless_vars: int, max_split: UInt[Array, ' p']
1473) -> None | UInt[Array, ' q']:
1474 """Initialize the `blocked_vars` field."""
1475 if filter_splitless_vars:
1476 (p,) = max_split.shape
1477 (blocked_vars,) = jnp.nonzero(
1478 max_split == 0, size=filter_splitless_vars, fill_value=p
1479 )
1480 return blocked_vars.astype(minimal_unsigned_dtype(p))
1481 # see `fully_used_variables` for the type cast
1482 else:
1483 return None
1486def add_chains(state: 'State', num_chains: int | None) -> 'State':
1487 """Extend chain-marked `_LazyArray` leaves to include a chain axis of size `num_chains`.
1489 Walks `state`, asks `chain_vmap_axes` where each leaf's chain axis lives,
1490 and wraps the carried `_LazyArray` so its factory creates the core array
1491 and then broadcasts a chain axis in at that position. To make
1492 `chain_vmap_axes` normalize against the chain-extended ``ndim``, the
1493 lookup is done on a shape preview built via `add_dummy_axis`. No-op when
1494 `num_chains` is `None`.
1496 Chain-marked leaves are required to be `_LazyArray` (or `None`); eager
1497 arrays at chain-marked positions are rejected so that all chain insertion
1498 happens at concretization time inside `shard_state`.
1499 """
1500 if num_chains is None:
1501 return state
1502 preview = add_dummy_axis(state)
1503 chain_axes = chain_vmap_axes(preview)
1505 def wrap(leaf: object, chain_axis: int | None) -> object:
1506 if chain_axis is None or leaf is None:
1507 return leaf
1508 assert isinstance(leaf, _LazyArray), (
1509 f'expected _LazyArray for chain-marked leaf, got {type(leaf).__name__}'
1510 )
1511 return _wrap_chain(leaf, chain_axis, num_chains)
1513 return tree.map(wrap, state, chain_axes, is_leaf=_is_lazy_or_none)
1516def parse_mesh(
1517 num_chains: int | None, mesh: Mesh | dict[str, int] | None
1518) -> Mesh | None:
1519 """Parse the `mesh` argument."""
1520 if mesh is None:
1521 return None
1523 # convert dict format to actual mesh
1524 if not isinstance(mesh, Mesh):
1525 assert set(mesh).issubset({'chains', 'data'})
1526 mesh = make_mesh(
1527 tuple(mesh.values()), tuple(mesh), axis_types=(AxisType.Auto,) * len(mesh)
1528 )
1530 # the chains mesh axis must be consistent with the number of chains
1531 if 'chains' in mesh.axis_names:
1532 if num_chains is None:
1533 msg = "mesh has a 'chains' axis but num_chains is None (scalar, no chain axis)"
1534 raise ValueError(msg)
1535 chains_axis = get_axis_size(mesh, 'chains')
1536 if num_chains % chains_axis:
1537 msg = (
1538 f"mesh 'chains' axis of size {chains_axis} does not divide "
1539 f'num_chains={num_chains}'
1540 )
1541 raise ValueError(msg)
1543 # check the axes we use are in auto mode
1544 assert 'chains' not in mesh.axis_names or 'chains' in mesh.auto_axes
1545 assert 'data' not in mesh.axis_names or 'data' in mesh.auto_axes
1547 return mesh
1550@partial(filter_jit, donate='all')
1551# jit and donate because otherwise type conversion would create copies
1552def remove_weak_types(x: PyTree[Array, 'T']) -> PyTree[Array, 'T']:
1553 """Make all types strong.
1555 This is to avoid recompilation in `run_mcmc` or `step`.
1556 """
1558 def remove_weak(x: T) -> T:
1559 if isinstance(x, Array) and x.weak_type:
1560 return cast(T, x.astype(x.dtype))
1561 else:
1562 return x
1564 return tree.map(remove_weak, x)
1567def shard_state(state: State) -> State:
1568 """Place all arrays on the appropriate devices, and instantiate lazily defined arrays."""
1569 mesh = state.config.mesh
1570 shard_leaf_mesh = partial(shard_leaf, mesh=mesh)
1571 return tree.map(
1572 shard_leaf_mesh,
1573 state,
1574 chain_vmap_axes(state),
1575 data_vmap_axes(state),
1576 is_leaf=lambda x: x is None or isinstance(x, _LazyArray),
1577 )
1580def leaf_partition_spec(
1581 ndim: int, chain_axis: int | None, data_axis: int | None, mesh: Mesh
1582) -> PartitionSpec:
1583 """Build a `PartitionSpec` for a leaf with the given chain/data axes."""
1584 spec = [None] * ndim
1585 if chain_axis is not None and 'chains' in mesh.axis_names:
1586 spec[chain_axis] = 'chains'
1587 if data_axis is not None and 'data' in mesh.axis_names:
1588 spec[data_axis] = 'data'
1590 # remove trailing Nones to be consistent with jax's output, it's useful
1591 # for comparing shardings during debugging
1592 while spec and spec[-1] is None:
1593 spec.pop()
1595 return PartitionSpec(*spec)
1598def shard_leaf(
1599 x: Shaped[Array, '*shape'] | None | Shaped[_LazyArray, '*shape'],
1600 chain_axis: int | None,
1601 data_axis: int | None,
1602 mesh: Mesh | None,
1603) -> Shaped[Array, '*shape'] | None:
1604 """Create `x` if it's lazy and shard it."""
1605 if x is None:
1606 return None
1608 if mesh is None:
1609 sharding = None
1610 else:
1611 spec = leaf_partition_spec(x.ndim, chain_axis, data_axis, mesh)
1612 sharding = NamedSharding(mesh, spec)
1614 if isinstance(x, _LazyArray):
1615 x = concretize_lazy_array(x, sharding)
1616 elif sharding is not None:
1617 x = device_put(x, sharding, donate=True)
1619 return x
1622@filter_jit
1623# jit such that in recent jax versions the shards are created on the right
1624# devices immediately instead of being created on the wrong device and then
1625# copied
1626def concretize_lazy_array(
1627 x: Shaped[_LazyArray, '*shape'], sharding: NamedSharding | None
1628) -> Shaped[Array, '*shape']:
1629 """Create an array from an abstract spec on the appropriate devices."""
1630 x = x()
1631 if sharding is not None:
1632 x = lax.with_sharding_constraint(x, sharding)
1633 return x
1636def all_none_or_not_none(*args: object) -> bool:
1637 is_none = [x is None for x in args]
1638 return all(is_none) or not any(is_none)
1641def asarray_or_none(x: object) -> Shaped[Array, '...'] | None:
1642 if x is None:
1643 return None
1644 return jnp.asarray(x)
1647class ReductionConfigs(TypedDict):
1648 """Fields of `StepConfig` related to reductions."""
1650 resid_reduction_config: ReductionConfig
1651 count_reduction_config: ReductionConfig
1652 prec_reduction_config: ReductionConfig
1653 prec_count_num_trees: int | None
1656def parse_reduction_configs(
1657 resid_reduction_config: ReductionConfig,
1658 count_reduction_config: ReductionConfig,
1659 prec_reduction_config: ReductionConfig,
1660 prec_count_num_trees: int | None | Literal['auto'],
1661 y: Float32[Array, ' n'] | Float32[Array, ' k n'] | Bool[Array, ' n'],
1662 num_trees: int,
1663 num_chains: int | None,
1664 mesh: Mesh | None,
1665) -> ReductionConfigs:
1666 """Determine settings for indexed reduces."""
1667 n = y.shape[-1]
1668 n //= get_axis_size(mesh, 'data') # per-device datapoints
1670 # chains are vmapped together on each device, so they share the per-step
1671 # memory of the per-tree reduction
1672 chains_per_device = (num_chains or 1) // get_axis_size(mesh, 'chains')
1674 # the reduction configs carry their own datapoint-batch settings (resolved
1675 # per-platform at run time when 'auto', see `ReductionConfig`), so they are
1676 # stored verbatim; only `prec_count_num_trees`, which does not depend on the
1677 # platform, is resolved here
1678 return dict(
1679 resid_reduction_config=resid_reduction_config,
1680 count_reduction_config=count_reduction_config,
1681 prec_reduction_config=prec_reduction_config,
1682 prec_count_num_trees=parse_prec_count_num_trees(
1683 prec_count_num_trees, num_trees, n * chains_per_device
1684 ),
1685 )
1688def parse_prec_count_num_trees(
1689 prec_count_num_trees: int | None | Literal['auto'], num_trees: int, n: int
1690) -> int | None:
1691 """Return the number of trees to process at a time or determine it automatically."""
1692 if prec_count_num_trees != 'auto':
1693 return prec_count_num_trees
1694 max_n_by_ntree = 2**27 # about 100M
1695 pcnt = max_n_by_ntree // max(1, n)
1696 pcnt = min(num_trees, pcnt)
1697 pcnt = max(1, pcnt)
1698 pcnt = search_divisor(
1699 pcnt, num_trees, max(1, pcnt // 2), max(1, min(num_trees, pcnt * 2))
1700 )
1701 if pcnt >= num_trees: 1701 ↛ 1703line 1701 didn't jump to line 1703 because the condition on line 1701 was always true
1702 pcnt = None
1703 return pcnt
1706def search_divisor(target_divisor: int, dividend: int, low: int, up: int) -> int:
1707 """Find the divisor closest to `target_divisor` in [low, up] if `target_divisor` is not already.
1709 If there is none, give up and return `target_divisor`.
1710 """
1711 assert target_divisor >= 1
1712 assert 1 <= low <= up <= dividend
1713 if dividend % target_divisor == 0:
1714 return target_divisor
1715 candidates = numpy.arange(low, up + 1)
1716 divisors = candidates[dividend % candidates == 0]
1717 if divisors.size == 0:
1718 return target_divisor
1719 penalty = numpy.abs(divisors - target_divisor)
1720 closest = numpy.argmin(penalty)
1721 return divisors[closest].item()
1724def get_axis_size(mesh: Mesh | None, axis_name: str) -> int:
1725 if mesh is None or axis_name not in mesh.shape:
1726 return 1
1727 else:
1728 return mesh.shape[axis_name]
1731def chol_with_gersh(
1732 mat: Float32[Array, '*batch_shape k k'], absolute_eps: bool = False
1733) -> Float32[Array, '*batch_shape k k']:
1734 """Cholesky with Gershgorin stabilization, supports batching."""
1735 return _chol_with_gersh_impl(mat, absolute_eps)
1738@partial(jnp.vectorize, signature='(k,k)->(k,k)', excluded=(1,))
1739def _chol_with_gersh_impl(
1740 mat: Float32[Array, '*batch_shape k k'], absolute_eps: bool
1741) -> Float32[Array, '*batch_shape k k']:
1742 # standardize to unit diagonal first, so the Gershgorin shift is relative to
1743 # each component's scale instead of a single absolute value set by the
1744 # largest one (which would swamp components with much smaller variance, e.g.
1745 # mixing a heavily-scaled continuous outcome with O(1) binary ones).
1746 # degenerate (non-positive or non-finite, e.g. infinite precision from a
1747 # constant outcome) diagonals fall back to the largest finite scale, which
1748 # keeps a diagonal matrix bit-identical to an unstandardized stabilization
1749 # and leaves an infinite diagonal infinite (pinning that leaf to zero)
1750 diag = jnp.diagonal(mat)
1751 finite_pos = jnp.isfinite(diag) & (diag > 0)
1752 ref = jnp.max(jnp.where(finite_pos, diag, 0.0), initial=0.0)
1753 scale = jnp.sqrt(jnp.where(finite_pos, diag, jnp.where(ref > 0, ref, 1.0)))
1754 mat = mat / (scale[:, None] * scale[None, :])
1755 rho = jnp.max(jnp.sum(jnp.abs(mat), axis=1), initial=0.0)
1756 eps = jnp.finfo(mat.dtype).eps
1757 u = mat.shape[0] * rho * eps
1758 if absolute_eps:
1759 u += eps
1760 mat = mat.at[jnp.diag_indices_from(mat)].add(u)
1761 return scale[:, None] * jnp.linalg.cholesky(mat)
1764def inv_via_chol_with_gersh(
1765 mat: Float32[Array, '*batch_shape k k'],
1766) -> Float32[Array, '*batch_shape k k']:
1767 """Compute matrix inverse via Cholesky with Gershgorin stabilization.
1769 DO NOT USE THIS FUNCTION UNLESS YOU REALLY NEED TO.
1770 """
1771 # mat = L L^T
1772 # mat^-1 = L^-T L^-1 = L^-T I L^-1 = L^-T (L^-T I)^T
1773 # I suspect this to be more accurate than (L^-1 I)^T (L^-1 I)
1774 L = chol_with_gersh(mat)
1775 eye = jnp.broadcast_to(jnp.eye(mat.shape[-1]), mat.shape)
1776 Ltinv = solve_triangular(L, eye, trans='T', lower=True)
1777 return solve_triangular(L, Ltinv.mT, trans='T', lower=True)
1780def split_key_for_chains(
1781 fun: Callable[[Key[Array, ''] | Key[Array, ' num_chains'], State], State],
1782) -> Callable[[Key[Array, ''], State], State]:
1783 """Split a single PRNG key into per-chain keys before calling `fun`.
1785 When the state is multichain, the input key is split into
1786 ``state.num_chains()`` keys. For single-chain states, the key is passed
1787 through unchanged.
1788 """
1790 @wraps(fun)
1791 def wrapped(key: Key[Array, ''], state: State) -> State:
1792 num_chains = state.num_chains()
1793 if num_chains is None:
1794 return fun(key, state)
1795 keys = random.split(key, num_chains)
1796 return fun(keys, state)
1798 return wrapped
1801def partition_specs(x: PyTree, mesh: Mesh) -> PyTree[PartitionSpec]:
1802 """Per-leaf `PartitionSpec`s derived from chain/data `field` markers.
1804 Each array leaf is sharded over ``'chains'`` along its chain axis and over
1805 ``'data'`` along its data axis, when those axes are marked (see `field`)
1806 and present in `mesh`; all other axes are replicated.
1808 Parameters
1809 ----------
1810 x
1811 A pytree of arrays carrying chain/data `field` markers.
1812 mesh
1813 The device mesh to shard over.
1815 Returns
1816 -------
1817 A pytree matching `x` with a `PartitionSpec` in place of each array leaf.
1818 """
1819 return tree.map(
1820 lambda leaf, ca, da: leaf_partition_spec(leaf.ndim, ca, da, mesh),
1821 x,
1822 chain_vmap_axes(x),
1823 data_vmap_axes(x),
1824 )
1827def shard_map_state(
1828 fun: Callable[[Key[Array, ''] | Key[Array, ' num_chains'], State], State],
1829) -> Callable[[Key[Array, ''] | Key[Array, ' num_chains'], State], State]:
1830 """Wrap a ``(keys, state) -> state`` function in a manual `jax.shard_map`.
1832 Uses `state.config.mesh` (static). No-op when the mesh is `None`. The keys
1833 input is sharded across ``'chains'`` when the state is multichain and
1834 ``'chains'`` is in the mesh; otherwise the keys are replicated. State
1835 leaves are sharded according to their `chains`/`data` field metadata. The
1836 output sharding matches the input sharding.
1837 """
1839 @wraps(fun)
1840 def wrapped(key: Key[Array, ''] | Key[Array, ' num_chains'], state: State) -> State:
1841 mesh = state.config.mesh
1842 if mesh is None:
1843 return fun(key, state)
1845 if state.has_chains and 'chains' in mesh.axis_names:
1846 key_spec = PartitionSpec('chains')
1847 else:
1848 key_spec = PartitionSpec()
1850 state_specs = partition_specs(state, mesh)
1852 mapped = shard_map(
1853 fun,
1854 mesh=mesh,
1855 in_specs=(key_spec, state_specs),
1856 out_specs=state_specs,
1857 **get_shard_map_patch_kwargs(),
1858 )
1859 return mapped(key, state)
1861 return wrapped
1864def vmap_chains(
1865 fun: Callable[[Key[Array, ''], State], State],
1866) -> Callable[[Key[Array, ' num_chains'] | Key[Array, ''], State], State]:
1867 """Vmap a ``(key, state) -> state`` function over chain axes.
1869 When the state is multichain, `keys` must have a leading chain axis and
1870 `fun` is vmapped over it together with the chain axes of `state`. For
1871 single-chain states, the function is called unchanged.
1872 """
1874 @wraps(fun)
1875 def wrapped(
1876 keys: Key[Array, ' num_chains'] | Key[Array, ''], state: State
1877 ) -> State:
1878 if not state.has_chains:
1879 return fun(keys, state)
1880 state_axes = chain_vmap_axes(state)
1881 vmapped_fun = vmap(fun, in_axes=(0, state_axes), out_axes=state_axes)
1882 return vmapped_fun(keys, state)
1884 return wrapped
1887class ShardMapPatchKwargs(TypedDict, total=False):
1888 check_vma: bool
1891def get_shard_map_patch_kwargs() -> ShardMapPatchKwargs:
1892 # bug: jax 0.8.1-0.8.2: vmap(shard_map(psum)), jax#34249; the
1893 # jax_disable_vmap_shmap_error config did not work.
1895 # WORKAROUND(jax<=0.8.2): remove this whole function when jax > 0.8.2
1896 buggy = ('0.8.1', '0.8.2')
1897 if jax.__version__ in buggy: 1897 ↛ 1898line 1897 didn't jump to line 1898 because the condition on line 1897 was never true
1898 return {'check_vma': False}
1899 return {}