Coverage for src/bartz/mcmcstep/_state.py: 99%
568 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-15 10:04 +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``.
106 Each marginal error variance (diagonal of the inverse) is inverse-gamma
107 with ``alpha = inv_wishart_marginal_nu / 2``, ``beta = rate_ii / 2``.
109 Set `nu` and `rate` to `None` to represent a precision held fixed at `value`
110 with no prior (e.g. the identity in binary regression).
111 """
113 nu: Float32[Array, ''] | None
114 """Degrees of freedom of the Wishart prior, or `None` if there is no prior."""
116 rate: Float32[Array, ''] | Float32[Array, 'k k'] | None
117 """The rate matrix of the Wishart prior (scalar for univariate), or `None`
118 if there is no prior. Equal to the inverse-gamma ``scale`` in the
119 univariate case."""
121 value: Float32[Array, '*chains k k'] | Float32[Array, '*chains'] = field(
122 chains=CHAIN_AXIS
123 )
124 """The precision matrix (scalar for univariate)."""
126 def __init__(
127 self,
128 nu: FloatLike | None,
129 rate: FloatLike | Float[ArrayLike, 'k k'] | None,
130 value: FloatLike
131 | Float[ArrayLike, '*chains k k']
132 | Float[ArrayLike, '*chains'],
133 ) -> None:
134 # `init` passes a deferred `_LazyArray` (cast to `Array`) for `value` to
135 # route it through sharding.
136 assert (nu is None) == (rate is None), 'set both or neither of nu and rate'
137 self.nu = None if nu is None else jnp.asarray(nu, jnp.float32)
138 self.rate = None if rate is None else jnp.asarray(rate, jnp.float32)
139 if isinstance(value, _LazyArray):
140 self.value = cast(Array, value)
141 else:
142 self.value = jnp.asarray(value, jnp.float32)
144 @classmethod
145 def from_inv_wishart_marginal_nu(
146 cls,
147 marginal_nu: FloatLike,
148 rate: FloatLike | Float[ArrayLike, 'k k'],
149 value: FloatLike
150 | Float[ArrayLike, '*chains k k']
151 | Float[ArrayLike, '*chains'],
152 ) -> 'Wishart':
153 """Build the prior from the degrees of freedom of each marginal variance.
155 Sets ``nu = marginal_nu + k - 1``, the inverse of
156 `inv_wishart_marginal_nu`.
157 """
158 if jnp.ndim(rate) == 0:
159 nu = marginal_nu
160 else:
161 k, _ = jnp.shape(rate)
162 nu = marginal_nu + (k - 1)
163 return cls(nu=nu, rate=rate, value=value)
165 @property
166 def inv_wishart_marginal_nu(self) -> Float32[Array, ''] | None:
167 """Degrees of freedom of the inverse-gamma marginal of each variance.
169 Equal to ``nu - k + 1``, or `None` if there is no prior.
170 """
171 if self.nu is None:
172 return None
173 else:
174 assert self.rate is not None
175 if self.rate.ndim == 0:
176 return self.nu
177 else:
178 k, _ = self.rate.shape
179 return self.nu - (k - 1)
182class DiagWishart(Wishart):
183 """A diagonal precision matrix with independent chi-square diagonal entries.
185 Despite the name this is not a Wishart restricted to diagonal matrices, but
186 a convenience type: a diagonal precision whose entries are mutually
187 independent, each with its own Gamma (scaled chi-square) prior. Only the
188 multivariate (matrix) case is supported.
190 With the same `nu` and `rate`, each component's variance has the same prior
191 as the corresponding marginal variance under the dense `Wishart`: the gamma
192 shape of each precision entry is ``inv_wishart_marginal_nu / 2``, not
193 ``nu / 2``.
195 A component with `rate` 0 has no prior; its precision is held fixed at its
196 `value` (1 for the binary components of a mixed regression).
198 Used for mixed binary-continuous regression and for continuous multivariate
199 regression with per-datapoint missingness.
200 """
202 def __init__(
203 self,
204 nu: FloatLike | None,
205 rate: FloatLike | Float[ArrayLike, 'k k'] | None,
206 value: FloatLike
207 | Float[ArrayLike, '*chains k k']
208 | Float[ArrayLike, '*chains'],
209 ) -> None:
210 # explicit (delegating) init so the static checker uses this signature
211 # instead of synthesizing a stricter one from the inherited fields
212 assert rate is None or jnp.ndim(rate) == 2, (
213 'DiagWishart supports only the multivariate (matrix) case'
214 )
215 super().__init__(nu, rate, value)
218class Forest(Module):
219 """Represents the MCMC state of a sum of trees."""
221 # Implementation note related to runtime type checking: Heap-array fields
222 # follow the `bartz.grove.TreesTrace` convention: the union-free integer
223 # trees are declared before `leaf_tree` and carry the bindable
224 # `half_tree_size` axis, while `leaf_tree` (and `p_nonterminal`) are
225 # checked against `2*half_tree_size`. Declaring a union-free `*chains`
226 # field first binds the variadic chain axis (plus `num_trees` and
227 # `half_tree_size`) before `leaf_tree`'s `... | ... k ...` union is
228 # evaluated, so the runtime typechecker can't mis-bind `*chains` against
229 # the `k` axis of a multivariate forest (the layouts are otherwise
230 # rank-ambiguous). No dummy anchor field needed.
232 var_tree: UInt[Array, '*chains num_trees half_tree_size'] = field(chains=CHAIN_AXIS)
233 """Variables/predictors/axes of decision rules."""
235 split_tree: UInt[Array, '*chains num_trees half_tree_size'] = field(
236 chains=CHAIN_AXIS
237 )
238 """Cutpoints/boundaries of decision rules."""
240 affluence_tree: Bool[Array, '*chains num_trees half_tree_size'] = field(
241 chains=CHAIN_AXIS
242 )
243 """Marks leaves that can be grown."""
245 leaf_tree: (
246 Float[Array, '*chains num_trees 2*half_tree_size']
247 | Float[Array, '*chains num_trees k 2*half_tree_size']
248 ) = field(chains=CHAIN_AXIS)
249 """The leaf values, in units of `leaf_unit`. The function computed by the
250 forest is ``offset + leaf_unit * (sum of leaves)``."""
252 leaf_unit: Float32[Array, ''] | Float32[Array, ' k']
253 """The storage unit of the leaves. Keeps the stored values O(1) whatever
254 the data units, so they do not under/overflow narrow `leaf_tree` dtypes.
255 Set to the marginal prior standard deviation of a leaf, rounded to a power
256 of two so converting to and from data units is exact."""
258 offset: Float32[Array, ''] | Float32[Array, ' k']
259 """Constant shift added to the scaled sum of trees, see `leaf_tree`."""
261 grow_prop_count: Int32[Array, '*chains'] = field(chains=CHAIN_AXIS)
262 """The number of grow proposals made during the last MCMC step."""
264 prune_prop_count: Int32[Array, '*chains'] = field(chains=CHAIN_AXIS)
265 """The number of prune proposals made during the last MCMC step."""
267 grow_acc_count: Int32[Array, '*chains'] = field(chains=CHAIN_AXIS)
268 """The number of grow moves accepted during the last MCMC step."""
270 prune_acc_count: Int32[Array, '*chains'] = field(chains=CHAIN_AXIS)
271 """The number of prune moves accepted during the last MCMC step."""
273 max_split: UInt[Array, ' p']
274 """The maximum split index for each predictor."""
276 blocked_vars: UInt[Array, ' q'] | None
277 """Indices of not to be used variables/predictors. This shall include at
278 least all the `i` that yield ``max_split[i] == 0``, otherwise behavior is
279 undefined."""
281 p_nonterminal: Float32[Array, ' 2*half_tree_size']
282 """The prior probability of each node being nonterminal (conditional on its
283 ancestors leaving at least one available decision rule). Includes the nodes
284 at maximum depth which shall be set to 0."""
286 p_propose_grow: Float32[Array, ' half_tree_size']
287 """The unnormalized probability of picking a leaf for a grow proposal."""
289 leaf_indices: UInt[Array, '*chains num_trees n'] = field(chains=CHAIN_AXIS, data=-1)
290 """The index of the leaf each datapoint falls into, for each tree, in the
291 largest version of the tree compatible with the last moves.
293 A pending prune (accepted prune or rejected grow, marked per-tree by
294 `to_prune`) is not yet applied to the indices; `step` folds it in at the
295 beginning of the next iteration. Evaluating the trees at these indices is
296 correct anyway because `leaf_tree` mirrors the value of a pruned node onto
297 its dangling children."""
299 to_prune: Bool[Array, '*chains num_trees'] = field(chains=CHAIN_AXIS)
300 """Whether the last move on each tree ended in a prune (accepted prune or
301 rejected grow) whose application to `leaf_indices` is still pending."""
303 move_node: Int32[Array, '*chains num_trees'] = field(chains=CHAIN_AXIS)
304 """The node the last move on each tree operated on (the leaf to grow or
305 the node to prune). Meaningful only where `to_prune` is set."""
307 count_tree: UInt32[Array, '*chains num_trees 2*half_tree_size'] | None = field(
308 chains=CHAIN_AXIS
309 )
310 """The number of datapoints per leaf. Valid at the leaves and at the nodes
311 involved in the latest moves, dirty elsewhere. `None` if there are
312 per-datapoint error scales and no minimum-points-per-node constraints,
313 which makes the counts unused."""
315 prec_tree: (
316 Float32[Array, '*chains num_trees 2*half_tree_size']
317 | Float32[Array, '*chains num_trees k k 2*half_tree_size']
318 | None
319 ) = field(chains=CHAIN_AXIS)
320 """The sum of `State.prec_scale` over the datapoints in each leaf, in the
321 same units; valid/dirty like `count_tree`. `None` if there are no
322 per-datapoint error scales, in which case `count_tree` takes its place."""
324 min_points_per_decision_node: Int32[Array, ''] | None
325 """The minimum number of data points in a decision node."""
327 min_points_per_leaf: Int32[Array, ''] | None
328 """The minimum number of data points in a leaf node."""
330 log_trans_prior: Float32[Array, '*chains num_trees'] | None = field(
331 chains=CHAIN_AXIS
332 )
333 """The log transition and prior Metropolis-Hastings ratio for the proposed
334 move on each tree."""
336 log_likelihood: Float32[Array, '*chains num_trees'] | None = field(
337 chains=CHAIN_AXIS
338 )
339 """The log likelihood ratio."""
341 leaf_prior_cov_inv: Wishart
342 """The prior precision matrix of a leaf, conditional on the tree structure
343 (a scalar inverse variance for univariate), with its Wishart prior. The
344 prior mean of a leaf is zero; the prior covariance of the sum of trees is
345 ``num_trees * leaf_prior_cov_inv.value^-1``. If `Wishart.nu` is `None`,
346 the precision is held fixed."""
348 log_s: Float32[Array, '*chains p'] | None = field(chains=CHAIN_AXIS)
349 """The logarithm of the prior probability for choosing a variable to split
350 along in a decision rule. Not normalized. Variables that do not have any
351 available decision rule at a given node are masked away. If `None`, use a
352 uniform distribution."""
354 theta: Float32[Array, '*chains'] | None = field(chains=CHAIN_AXIS)
355 """The concentration parameter for the Dirichlet prior on the variable
356 distribution `s`. If not set, `log_s` is left constant."""
358 a: Float32[Array, ''] | None
359 """Parameter of the prior on `theta`. If not set, `theta` is left constant."""
361 b: Float32[Array, ''] | None
362 """Parameter of the prior on `theta`. If not set, `theta` is left constant."""
364 rho: Float32[Array, ''] | None
365 """Parameter of the prior on `theta`. If not set, `theta` is left constant."""
367 @property
368 def has_chains(self) -> bool:
369 """Whether non-constant attributes in the forest carry an explicit chain axis."""
370 return self.var_tree.ndim > 2
373class StepConfig(Module):
374 """Options for the MCMC step."""
376 steps_done: Int32[Array, '']
377 """The number of MCMC steps completed so far."""
379 sparse_on_at: Int32[Array, ''] | None
380 """After how many steps to turn on variable selection. If `None`, variable
381 selection is disabled."""
383 resid_reduction_config: ReductionConfig
384 """How to sum the residuals in each leaf."""
386 count_reduction_config: ReductionConfig
387 """How to count the datapoints in each leaf."""
389 prec_reduction_config: ReductionConfig
390 """How to sum the likelihood precisions in each leaf."""
392 prec_count_num_trees: int | None = field(static=True)
393 """Batch size for processing trees to compute count and prec trees."""
395 sequential_unroll: int | bool = field(static=True)
396 """How much to unroll the sequential accept/reject loop over trees in
397 `step`. See the ``unroll`` argument of `jax.lax.scan`."""
399 augment: bool = field(static=True)
400 """Whether to account exactly, via data augmentation, for the decision
401 rules forbidden by the ancestors of each node when updating
402 `Forest.log_s`."""
404 mesh: Mesh | None = field(static=True)
405 """The mesh used to shard data and computation across multiple devices."""
407 leaf_quantization: Int32[Array, ''] | None = None
408 """If set, quantize the leaves to multiples of ``eps(resid dtype) * 2 **
409 leaf_quantization`` in `State.resid_eff_scale` units, which makes the
410 running updates of `State.resid` mostly exact, (almost) stopping their
411 random-walk rounding drift, assuming ``|resid| < 2 ** (leaf_quantization +
412 1)`` holds (in the same units) for most datapoints most of the time.
413 Intended mostly for use with float16 residuals. Sensible settings are 0 and
414 1, with 0 leaving some drift, 1 practically no drift, and no setting above
415 1 justifying the reduced accuracy. With enough datapoints or trees the
416 MCMC breaks down because the sampled leaf variation becomes smaller than
417 the quantum, so this setting can not be used liberally."""
419 @property
420 def data_sharded(self) -> bool:
421 """Whether the data axis is sharded across devices."""
422 return self.mesh is not None and 'data' in self.mesh.axis_names
425class State(Module):
426 """Represents the MCMC state of BART."""
428 _chain_anchor: Float32[Array, '*chains'] = field(chains=CHAIN_AXIS)
429 """Unused per-chain scalar, declared first as a runtime-typechecker anchor.
430 Its single (union-free) ``*chains`` annotation binds the variadic chain
431 axis before the ``... | ... k ...`` unions of `z`/`resid` (z over the
432 binary-outcome ``kb`` axis) are checked; otherwise those can mis-bind
433 ``*chains`` against the outcome axis for a multivariate-without-chains
434 state (the layouts are rank-ambiguous). Unlike `Forest`, `State` has no
435 genuine union-free chain field to reorder into this slot, so a dummy one is
436 carried."""
438 X: UInt[Array, 'p n'] = field(data=-1)
439 """The predictors."""
441 y: Float32[Array, ' n'] | Float32[Array, 'k n'] = field(data=-1)
442 """The response, in data units. Binary components are stored as 0/1.
443 Missing values are replaced by `Forest.offset`."""
445 z: None | Float32[Array, '*chains n'] | Float32[Array, '*chains kb n'] = field(
446 chains=CHAIN_AXIS, data=-1
447 )
448 """The latent outcomes for binary regression. `None` in continuous
449 regression. In the mixed binary-continuous case, only the binary outcome
450 components are stored."""
452 binary_indices: None | Int32[Array, ' kb']
453 """The indices of binary outcome components in the full list of outcome
454 components. `None` when there are no binary components."""
456 resid: Float[Array, '*chains n'] | Float[Array, '*chains k n'] = field(
457 chains=CHAIN_AXIS, data=-1
458 )
459 """The residuals, ``resid_unit * resid = (y or z) - sum of trees``."""
461 resid_unit: Float32[Array, ''] | Float32[Array, ' k']
462 """The storage unit of `resid` (see `init`'s ``resid_dtype``), same scheme
463 as `Forest.leaf_unit`. Equal to ``leaf_unit * sqrt(num_trees)`` (the
464 marginal prior standard deviation of the sum of trees) rounded to a power
465 of two."""
467 resid_eff_scale: Float32[Array, '*chains'] | Float32[Array, '*chains k'] = field(
468 chains=CHAIN_AXIS
469 )
470 """The measured scale of the residuals (their precision-weighted root mean
471 square in data units), rounded to a power of two. Sets the leaf
472 quantization grid (see `StepConfig.leaf_quantization`). Initialized to
473 `resid_unit`, then tracks the MCMC (while the storage unit of `resid` stays
474 fixed at `resid_unit`). Not updated in purely binary regression, where the
475 scale of the latent residuals should be stable."""
477 resid_inexact_integral: Float32[Array, '*chains'] | Float32[Array, '*chains k'] = (
478 field(chains=CHAIN_AXIS)
479 )
480 """Sum over the MCMC steps done of the mean square of the residuals (in
481 `resid_unit` units) large enough that their running updates round. Used by
482 `sum_trees_eps` to estimate the accumulated rounding drift."""
484 error_cov_inv: Wishart
485 """The inverse error covariance (``error_cov_inv.value``, scalar for
486 univariate) with its Wishart prior. Fixed at the identity with no prior in
487 binary regression."""
489 error_scale: Float32[Array, ' n'] | Float32[Array, 'k n'] | None = field(data=-1)
490 """The per-datapoint error scales (the ``error_scale`` argument of `init`).
491 The error precision on a datapoint is ``error_cov_inv.value /
492 outer(error_scale, error_scale)``. For binary components the (fixed, unit)
493 probit latent error is scaled instead, so the success probability is
494 ``Phi(sum of trees / error_scale)``. `inv_sdev_scale` and `prec_scale` are
495 derived from this and the missingness mask."""
497 prec_scale: Float[Array, ' n'] | Float[Array, 'k k n'] | None = field(data=-1)
498 """The scale on the error precision, ``prec_scale = outer(inv_sdev_scale,
499 inv_sdev_scale)`` per datapoint (``inv_sdev_scale ** 2`` for scalar
500 scales), so it's in units of ``inv_sdev_unit ** 2``. Stored, like
501 `inv_sdev_scale`, in `init`'s ``prec_scale_dtype``."""
503 inv_sdev_scale: Float[Array, ' n'] | Float[Array, 'k n'] | None = field(data=-1)
504 """``inv_sdev_scale * inv_sdev_unit = 1 / error_scale``, zeroed at missing
505 datapoints. Not `None` when ``missing`` is set even if fit without error
506 scales."""
508 inv_sdev_unit: Float32[Array, ''] | Float32[Array, ' k']
509 """The storage unit of `inv_sdev_scale`, to avoid under/overflow with short
510 dtypes; same scheme as `Forest.leaf_unit`. Set to the root mean square of
511 ``1 / error_scale`` over non-missing datapoints, rounded to a power of two;
512 1 if fit without error scales. Constant along the MCMC."""
514 n_non_missing: Int32[Array, ''] | Int32[Array, ' k']
515 """The number of non-missing datapoints."""
517 sum_diag_prec_scale: Float32[Array, ''] | Float32[Array, ' k']
518 """``sum(1 / error_scale ** 2)`` over non-missing datapoints; equal to
519 `n_non_missing` if fit without error scales."""
521 forest: Forest
522 """The sum of trees model."""
524 config: StepConfig
525 """Metadata and configurations for the MCMC step."""
527 @property
528 def has_chains(self) -> bool:
529 """Whether non-constant attributes in the state carry an explicit chain axis."""
530 return self.forest.has_chains
532 def num_chains(self) -> int | None:
533 """Return the number of chains, or `None` if not multichain."""
534 if not self.has_chains:
535 return None
536 c = chain_vmap_axes(self.forest).var_tree
537 return self.forest.var_tree.shape[c]
539 @jit
540 def sum_trees_eps(self) -> Float32[Array, ''] | Float32[Array, ' k']:
541 """Estimate the absolute accuracy limit of the sum of trees (in data units).
543 The analogue of ``finfo(dtype).eps`` for the sum of trees. This combines three
544 terms: floating point resolution, random walk accumulation of numerical error
545 on running residuals, and if leaf quantization is active, breakdown of the
546 mcmc due to the quantization being too coarse. The latter term is currently
547 broken actually, sorry about that.
548 """
549 resolution, drift, snap = self._sum_trees_eps()
550 return jnp.maximum(jnp.maximum(resolution, drift), snap)
552 def _sum_trees_eps(
553 self,
554 ) -> tuple[
555 Float32[Array, ''] | Float32[Array, ' k'],
556 Float32[Array, ''] | Float32[Array, ' k'],
557 Float32[Array, ''] | Float32[Array, ' k'],
558 ]:
559 """Return the resolution, drift, and snap terms of `sum_trees_eps` separately."""
560 eps_leaf = jnp.finfo(self.forest.leaf_tree.dtype).eps
561 eps_resid = jnp.finfo(self.resid.dtype).eps
563 # rounding of the stored leaves at their typical magnitude `leaf_unit`
564 dtype_quantum = eps_leaf * self.forest.leaf_unit
566 if self.config.leaf_quantization is None:
567 resolution = dtype_quantum
568 # float leaf storage rounds relative to the leaf magnitude, so
569 # leaves near zero keep a fine grid and sampling is not pinned
570 snap = jnp.zeros_like(dtype_quantum)
571 else:
572 # quantized leaves all sit on one grid closed under addition (the
573 # scales are powers of two, so leaf storage rounds within the
574 # grid even where its spacing is coarser), so the sum resolves
575 # multiples of the quantum whatever the number of trees
576 q = self.config.leaf_quantization
577 scale_chain_axes = range(self.resid_eff_scale.ndim - self.resid_unit.ndim)
578 eff_scale = jnp.mean(self.resid_eff_scale, axis=tuple(scale_chain_axes))
579 managed_quantum = eps_resid * eff_scale * 2.0**q
580 resolution = jnp.maximum(managed_quantum, dtype_quantum)
582 # determine average error precision; the folded `inv_sdev_unit`
583 # puts the stored `prec_scale` mean in data units
584 prec = scaled_error_cov_inv(self)
585 n_eff = jnp.maximum(self.n_non_missing, 1)
586 if self.prec_scale is not None:
587 prec *= self.prec_scale.sum(axis=-1, dtype=jnp.float32) / n_eff
589 # convert precision to variance and average it over chains
590 if self.resid_unit.ndim:
591 error_var = jnp.diagonal(
592 inv_via_chol_with_gersh(prec), axis1=-2, axis2=-1
593 )
594 else:
595 error_var = jnp.reciprocal(prec)
596 var_chain_axes = range(error_var.ndim - self.resid_unit.ndim)
597 error_var = jnp.mean(error_var, axis=tuple(var_chain_axes))
599 # a quantized leaf moves only when its full conditional mean
600 # crosses half a quantum; the mean responds to the leaf's average
601 # residual through the posterior shrinkage factor s = t / (1 + t),
602 # t = (leaf_unit / error_sdev)^2 * n_leaf, so residual features
603 # smaller than quantum / (2 s) cannot move the sampler. Typical
604 # values are used for the error variance (chain mean, datapoint
605 # mean of prec_scale) and n_leaf (datapoints over leaves per tree).
606 leaves_per_tree = 1.0 + jnp.count_nonzero(self.forest.split_tree, axis=-1)
607 n_leaf = n_eff / jnp.mean(leaves_per_tree)
608 t = jnp.square(self.forest.leaf_unit) * n_leaf / error_var
609 snap = managed_quantum / 2.0 * (1.0 + jnp.reciprocal(t))
611 # random walk drift of the accumulated rounding errors: each of the
612 # `num_trees` updates per step charges eps_resid * |resid| to the
613 # residuals large enough to round, whose mean square is integrated
614 # over the steps done in `resid_inexact_integral`
615 integral = self.resid_inexact_integral
616 chain_axes = range(integral.ndim - self.resid_unit.ndim)
617 integral = jnp.mean(integral, axis=tuple(chain_axes))
618 *_, num_trees, _ = self.forest.var_tree.shape
619 drift = eps_resid * self.resid_unit * jnp.sqrt(num_trees * integral)
621 return resolution, drift, snap
624def check_diagonal(rate: Float32[Array, 'k k']) -> Float32[Array, 'k k']:
625 """Raise if the Wishart `rate` is not diagonal."""
626 diag = jnp.diag(jnp.diag(rate))
627 return error_if(rate, jnp.any(rate != diag), 'error_cov_inv.rate must be diagonal')
630def check_binary_unit_precision(
631 value: Float32[Array, 'k k'], binary_mask: Sequence[bool]
632) -> Float32[Array, 'k k']:
633 """Raise if the binary diagonal entries of `value` are not fixed at 1."""
634 binary = jnp.array(binary_mask)
635 off_unit = jnp.any(binary & (jnp.diag(value) != 1.0))
636 return error_if(
637 value,
638 off_unit,
639 'binary error precision must be 1 (the default for a zero rate)',
640 )
643def init_shape_shifting_parameters(
644 *,
645 y: Float32[Array, ' n'] | Float32[Array, 'k n'],
646 outcome_type: OutcomeType | list[OutcomeType],
647 offset: Float32[Array, ''] | Float32[Array, ' k'],
648 error_scale: Float32[ArrayLike, ' n'] | Float32[ArrayLike, 'k n'] | None,
649 error_cov_inv: Wishart | None,
650 leaf_prior_cov_inv: Wishart,
651 missing: Bool[ArrayLike, ' n'] | Bool[ArrayLike, 'k n'] | None,
652) -> tuple[bool, tuple[int, ...], Wishart, None | Int32[Array, ' kb']]:
653 """
654 Check and initialize parameters that change array type/shape based on outcome kind.
656 Parameters
657 ----------
658 y
659 The response variable (used only for shape checks).
660 outcome_type
661 Whether the regression is continuous or binary. Can be a list of
662 `OutcomeType` for per-component specification in the multivariate case.
663 offset
664 The offset to add to the predictions.
665 error_scale
666 Per-datapoint error scales, ``(n,)`` or ``(k, n)`` (used only for
667 shape checks).
668 error_cov_inv
669 The Wishart prior on the error precision and its initial value, or
670 `None` for binary regression. The mixed and partial-missing diagonal
671 modes require a `DiagWishart`; in the mixed case the binary components
672 must have an initial precision of 1.
673 leaf_prior_cov_inv
674 The Wishart prior on the leaf precision and its initial value (used
675 only for shape and type checks).
676 missing
677 The per-datapoint missingness mask, used to detect partial missingness
678 (2-D mask) so that diagonal-mode initialization is selected.
680 Returns
681 -------
682 is_binary
683 Whether all outcomes are binary.
684 kshape
685 The outcome shape, empty for univariate, (k,) for multivariate.
686 error_cov_inv
687 The Wishart prior with its initial value resolved for the outcome kind.
688 binary_indices
689 The indices of binary outcome components, or `None` if there are none.
690 """
691 kshape = offset.shape
693 # determine per-component outcome kinds
694 if isinstance(outcome_type, list):
695 assert kshape, 'per-component outcome_type requires multivariate y'
696 (k,) = kshape
697 assert len(outcome_type) == k
698 binary_mask = [t is OutcomeType.binary for t in outcome_type]
699 is_binary = all(binary_mask)
700 is_mixed = any(binary_mask) and not is_binary
701 else:
702 is_binary = outcome_type is OutcomeType.binary
703 is_mixed = False
705 if is_mixed:
706 binary_indices = jnp.array([i for i, b in enumerate(binary_mask) if b])
707 else:
708 binary_indices = None
710 partial_missing = missing is not None and missing.ndim == 2 and kshape
712 assert (
713 error_scale is None
714 or error_scale.shape == y.shape # (k, n)
715 or error_scale.shape == y.shape[-1:] # (n,)
716 )
718 # All-binary: no prior, the precision is fixed at the identity.
719 if is_binary:
720 assert error_cov_inv is None, 'no error covariance prior in binary regression'
721 value = jnp.eye(kshape[0]) if kshape else jnp.array(1.0)
722 error_cov_inv = Wishart(nu=None, rate=None, value=value)
724 # Mixed binary-continuous, or continuous-mv with 2-D missingness: diagonal
725 # error covariance, updated component-wise. The caller must supply a
726 # `DiagWishart`; in the mixed case the binary components must have unit
727 # initial precision (see `DiagWishart`).
728 elif is_mixed or partial_missing:
729 assert isinstance(error_cov_inv, DiagWishart), (
730 'mixed binary-continuous or partial-missing regression requires a '
731 'DiagWishart error_cov_inv prior'
732 )
733 assert error_cov_inv.rate is not None
734 assert error_cov_inv.rate.shape == 2 * kshape
735 assert error_cov_inv.value.shape == 2 * kshape
736 rate = check_diagonal(error_cov_inv.rate)
737 value = check_diagonal(error_cov_inv.value)
738 if is_mixed:
739 value = check_binary_unit_precision(value, binary_mask)
740 error_cov_inv = replace(error_cov_inv, rate=rate, value=value)
742 # All-continuous: a dense `Wishart`.
743 else:
744 assert error_cov_inv is not None
745 assert type(error_cov_inv) is Wishart, (
746 'continuous regression requires a dense Wishart error_cov_inv prior'
747 )
748 rate = error_cov_inv.rate
749 assert rate is not None
750 assert rate.shape == 2 * kshape
751 assert error_cov_inv.value.shape == 2 * kshape
753 assert y.shape[:-1] == kshape
755 assert type(leaf_prior_cov_inv) is Wishart, (
756 'leaf_prior_cov_inv must be a dense Wishart'
757 )
758 assert leaf_prior_cov_inv.value.shape == 2 * kshape
759 if leaf_prior_cov_inv.rate is not None:
760 assert leaf_prior_cov_inv.rate.shape == 2 * kshape
762 return is_binary, kshape, error_cov_inv, binary_indices
765def check_splitless_vars(
766 filter_splitless_vars: int,
767 max_split: UInt[Array, ' p'],
768 offset: Float32[Array, ''] | Float32[Array, ' k'],
769) -> Float32[Array, ''] | Float32[Array, ' k']:
770 """Check there aren't too many deactivated predictors."""
771 msg = (
772 f'there are more than {filter_splitless_vars=} predictors with no splits, '
773 'please increase `filter_splitless_vars` or investigate the missing splits'
774 )
775 return error_if(offset, jnp.sum(max_split == 0) > filter_splitless_vars, msg)
778def parse_outcome_type(
779 outcome_type: 'OutcomeType | str | Sequence[OutcomeType | str]',
780) -> 'OutcomeType | list[OutcomeType]':
781 """Normalize outcome_type to enum (or list of enums)."""
782 if isinstance(outcome_type, Sequence) and not isinstance(outcome_type, str):
783 return [OutcomeType(t) for t in outcome_type]
784 else:
785 return OutcomeType(outcome_type)
788def parse_p_nonterminal(
789 p_nonterminal: Float32[ArrayLike, ' d_minus_1'],
790) -> Float32[Array, ' d_minus_1+1']:
791 """Check it's in (0, 1) and pad with a 0 at the end."""
792 p_nonterminal = jnp.asarray(p_nonterminal)
793 ok = (p_nonterminal > 0) & (p_nonterminal < 1)
794 p_nonterminal = error_if(p_nonterminal, ~ok, 'p_nonterminal must be in (0, 1)')
795 return jnp.pad(p_nonterminal, (0, 1))
798def make_p_nonterminal(
799 d: int, alpha: FloatLike = 0.95, beta: FloatLike = 2.0
800) -> Float32[Array, ' {d}-1']:
801 """Prepare the `p_nonterminal` argument to `init`.
803 It is calculated according to the formula:
805 P_nt(depth) = alpha / (1 + depth)^beta, with depth 0-based
807 Parameters
808 ----------
809 d
810 The maximum depth of the trees (d=1 means tree with only root node)
811 alpha
812 The a priori probability of the root node having children, conditional
813 on it being possible
814 beta
815 The exponent of the power decay of the probability of having children
816 with depth.
818 Returns
819 -------
820 An array of probabilities, one per tree level but the last.
821 """
822 assert d >= 1
823 depth = jnp.arange(d - 1)
824 return alpha / (1 + depth).astype(float) ** beta
827@float32_matmuls
828def init(
829 *,
830 X: UInt[ArrayLike, 'p n'],
831 y: Float32[ArrayLike, ' n'] | Float32[ArrayLike, ' k n'],
832 outcome_type: OutcomeType | str | Sequence[OutcomeType | str] = 'continuous',
833 offset: FloatLike | Float[ArrayLike, ' k'],
834 max_split: UInt[ArrayLike, ' p'],
835 num_trees: int,
836 p_nonterminal: Float32[ArrayLike, ' d_minus_1'],
837 leaf_prior_cov_inv: Wishart,
838 leaf_dtype: DTypeLike = jnp.float32,
839 prec_scale_dtype: DTypeLike = jnp.float32,
840 resid_dtype: DTypeLike = jnp.float32,
841 leaf_quantization: int | Integer[ArrayLike, ''] | None = None,
842 error_cov_inv: Wishart | None = None,
843 error_scale: Float32[ArrayLike, ' n'] | Float32[ArrayLike, 'k n'] | None = None,
844 missing: Bool[ArrayLike, ' n'] | Bool[ArrayLike, 'k n'] | None = None,
845 min_points_per_decision_node: int | Integer[ArrayLike, ''] | None = None,
846 min_points_per_leaf: int | Integer[ArrayLike, ''] | None = None,
847 resid_reduction_config: ReductionConfig = AutoBatchedReduction(),
848 count_reduction_config: ReductionConfig = AutoOneHotReduction(),
849 prec_reduction_config: ReductionConfig = AutoOneHotReduction(),
850 prec_count_num_trees: int | None | Literal['auto'] = 'auto',
851 sequential_unroll: int | bool = 2,
852 save_ratios: bool = False,
853 filter_splitless_vars: int = 0,
854 log_s: Float32[ArrayLike, ' p'] | None = None,
855 theta: FloatLike | None = None,
856 a: FloatLike | None = None,
857 b: FloatLike | None = None,
858 rho: FloatLike | None = None,
859 sparse_on_at: int | Integer[ArrayLike, ''] | None = None,
860 augment: bool = True,
861 num_chains: int | None = None,
862 mesh: Mesh | dict[str, int] | None = None,
863) -> State:
864 """
865 Make a BART posterior sampling MCMC initial state.
867 Parameters
868 ----------
869 X
870 The predictors. Note this is trasposed compared to the usual convention.
871 y
872 The response. If two-dimensional, the outcome is multivariate with the
873 first axis indicating the component. For binary data, non-zero means 1,
874 zero means 0.
875 outcome_type
876 Whether the regression is continuous or binary (probit). Can also be a
877 sequence of `OutcomeType` values, one per outcome component, for mixed
878 binary-continuous multivariate regression.
879 offset
880 Constant shift added to the sum of trees. 0 if not specified.
881 max_split
882 The maximum split index for each variable. All split ranges start at 1.
883 num_trees
884 The number of trees in the forest.
885 p_nonterminal
886 The probability of a nonterminal node at each depth. The maximum depth
887 of trees is fixed by the length of this array. Use `make_p_nonterminal`
888 to set it with the conventional formula.
889 leaf_prior_cov_inv
890 The Wishart prior on the leaf precision (a scalar inverse variance for
891 univariate) and its initial value; see `Forest.leaf_prior_cov_inv`.
892 Set `Wishart.nu` and `Wishart.rate` to `None` to hold the precision
893 fixed. Must be a dense `Wishart`, not a `DiagWishart`.
894 leaf_dtype
895 prec_scale_dtype
896 resid_dtype
897 Storage dtypes for, respectively: leaves (`Forest.leaf_tree`), derived
898 error scales (`State.prec_scale` and `State.inv_sdev_scale`; the raw
899 `State.error_scale` stays float32), and running residuals
900 (`State.resid`). These quantities are stored in scaled units so that
901 narrow dtypes do not under/overflow (see `Forest.leaf_unit`), though
902 float16 error scales may still overflow if their dynamic range is high
903 enough. Leaf full conditionals are always computed and sampled in
904 float32. Narrow residual dtypes may easily break the MCMC,
905 `resid_dtype` is an experimental setting.
906 leaf_quantization
907 Quantize the leaves to (almost) stop the numerical drift of the running
908 residuals, see `StepConfig.leaf_quantization`.
909 error_cov_inv
910 The prior and initial value of the error term precision matrix (see
911 `Wishart`). Leave it unspecified for binary regression. Mixed
912 binary-continuous and per-outcome-component missingness require a
913 `DiagWishart`.
914 error_scale
915 Per-datapoint error scales (called ``w`` in the R package BART3); see
916 `State.error_scale`. An unspecified `error_scale` is equivalent to
917 ``error_scale = 1`` for all datapoints.
918 missing
919 Boolean mask indicating which datapoints are missing. `True` marks
920 entries to be ignored by the MCMC. The masked values of `y` may be
921 anything, even non-finite. If `missing` is 2-D, `error_cov_inv` must
922 be a `DiagWishart`.
923 min_points_per_decision_node
924 min_points_per_leaf
925 The minimum number of datapoints in a decision node and in a leaf,
926 respectively; 0 if not specified. The leaf constraint is not taken into
927 account in the proposal distribution because that would be expensive.
928 The two are independent and not checked for coherence; it makes sense
929 to set ``min_points_per_decision_node >= 2 * min_points_per_leaf``.
930 resid_reduction_config
931 count_reduction_config
932 prec_reduction_config
933 How to sum the residuals, count the datapoints, and sum the likelihood
934 precisions in each leaf, respectively. See `ReductionConfig` and its
935 subclasses.
936 prec_count_num_trees
937 The number of trees to process at a time when counting datapoints or
938 computing the likelihood precision. If `None`, do all trees at once,
939 which may use too much memory on cpu. If 'auto' (default), it's chosen
940 automatically.
941 sequential_unroll
942 See `StepConfig.sequential_unroll`. Unrolling may speed up the MCMC at
943 the cost of longer compilation; 1 means no unrolling.
944 save_ratios
945 Whether to save the Metropolis-Hastings ratios.
946 filter_splitless_vars
947 The maximum number of variables without splits that can be ignored. If
948 there are more, `init` raises an exception.
949 log_s
950 theta
951 a
952 b
953 rho
954 Sparsity (variable selection) parameters, see `Forest.log_s` and
955 `Forest.theta`. If `rho`, `a`, `b` are set, an unspecified `theta` is
956 initialized to `rho`; if `theta` is set, an unspecified `log_s` is
957 initialized to uniform.
958 sparse_on_at
959 After how many MCMC steps to turn on variable selection.
960 augment
961 See `StepConfig.augment`. If disabled, forbidden decision rules are
962 ignored when counting variable usage, which may be faster but is
963 an approximation.
964 num_chains
965 The number of independent MCMC chains. Single chain with scalar values
966 if not specified.
967 mesh
968 A jax mesh used to shard data and computation across multiple devices.
969 If it has a 'chains' axis, that axis is used to shard the chains. If it
970 has a 'data' axis, that axis is used to shard the datapoints.
972 As a shorthand, if a dictionary mapping axis names to axis size is
973 passed, the corresponding mesh is created, e.g., ``dict(chains=4,
974 data=2)`` will let jax pick 8 devices to split chains (which must be a
975 multiple of 4) across 4 pairs of devices, where in each pair the data
976 is split in two.
978 Note: if a mesh is passed, the arrays are always sharded according to
979 it. In particular even if the mesh has no 'chains' or 'data' axis, the
980 arrays will be replicated on all devices in the mesh.
982 Returns
983 -------
984 An initialized BART MCMC state.
986 Raises
987 ------
988 ValueError
989 If arguments unused in binary regression are set.
991 Notes
992 -----
993 In decision nodes, the values in ``X[i, :]`` are compared to a cutpoint out
994 of the range ``[1, 2, ..., max_split[i]]``. A point belongs to the left
995 child iff ``X[i, j] < cutpoint``. Thus it makes sense for ``X[i, :]`` to be
996 integers in the range ``[0, 1, ..., max_split[i]]``.
998 In general the arrays passed to this function as arguments may be donated,
999 invalidating them. Create copies before passing them to `init` if this
1000 happens and you need them again.
1001 """
1002 # convert to array all array-like arguments that are used in other
1003 # configurations but don't need further processing themselves
1004 X = jnp.asarray(X)
1005 y = jnp.asarray(y)
1006 assert y.dtype == jnp.float32
1007 offset = jnp.asarray(offset)
1008 max_split = jnp.asarray(max_split)
1009 error_scale = asarray_or_none(error_scale)
1010 missing = asarray_or_none(missing)
1011 assert missing is None or missing.ndim <= y.ndim
1013 # normalize outcome_type to enum (or list of enums)
1014 outcome_type = parse_outcome_type(outcome_type)
1016 # check p_nonterminal and pad it with a 0 at the end (still not final shape)
1017 p_nonterminal = parse_p_nonterminal(p_nonterminal)
1019 # process arguments that change depending on outcome type
1020 is_binary, kshape, error_cov_inv, binary_indices = init_shape_shifting_parameters(
1021 y=y,
1022 outcome_type=outcome_type,
1023 offset=offset,
1024 error_scale=error_scale,
1025 error_cov_inv=error_cov_inv,
1026 leaf_prior_cov_inv=leaf_prior_cov_inv,
1027 missing=missing,
1028 )
1030 storage = determine_storage_params(
1031 leaf_dtype=leaf_dtype,
1032 prec_scale_dtype=prec_scale_dtype,
1033 resid_dtype=resid_dtype,
1034 leaf_prior_cov_inv=leaf_prior_cov_inv.value,
1035 kshape=kshape,
1036 num_trees=num_trees,
1037 )
1039 # extract array sizes from arguments
1040 (max_depth,) = p_nonterminal.shape
1041 p, n = X.shape
1043 # check and initialize sparsity parameters
1044 if not all_none_or_not_none(rho, a, b): 1044 ↛ 1045line 1044 didn't jump to line 1045 because the condition on line 1044 was never true
1045 msg = 'rho, a, b are not either all `None` or all set'
1046 raise ValueError(msg)
1047 if theta is None and rho is not None:
1048 theta = rho
1049 if log_s is None and theta is not None:
1050 log_s = jnp.zeros(max_split.size)
1051 if not all_none_or_not_none(theta, sparse_on_at): 1051 ↛ 1052line 1051 didn't jump to line 1052 because the condition on line 1051 was never true
1052 msg = 'sparsity params (either theta or rho,a,b) and sparse_on_at must be either all None or all set'
1053 raise ValueError(msg)
1055 # determine settings for reductions
1056 mesh = parse_mesh(num_chains, mesh)
1057 red_cfg = parse_reduction_configs(
1058 resid_reduction_config=resid_reduction_config,
1059 count_reduction_config=count_reduction_config,
1060 prec_reduction_config=prec_reduction_config,
1061 prec_count_num_trees=prec_count_num_trees,
1062 y=y,
1063 num_trees=num_trees,
1064 num_chains=num_chains,
1065 mesh=mesh,
1066 )
1068 # check there aren't too many deactivated predictors
1069 offset = check_splitless_vars(filter_splitless_vars, max_split, offset)
1071 tree_size = 2**max_depth
1073 # Assemble the state, shard it, then fill in the post-shard fields. This
1074 # whole region runs with type-checking disabled because the state carries
1075 # deliberately wrong-typed intermediates parked in its fields for sharding:
1076 # `_LazyArray` leaves (each chain-bearing leaf is built at its core no-chain
1077 # shape, then `add_chains` wraps it to broadcast in the chain axis) and the
1078 # user `missing` mask in the `inv_sdev_scale` slot. The context ends once
1079 # every field has been replaced by its final, correctly-typed array.
1080 with jaxtyping_disabled():
1081 state = State(
1082 _chain_anchor=lazy(jnp.zeros, ()), # typechecker chain anchor
1083 X=X,
1084 y=y,
1085 z=(
1086 lazy(jnp.full, y.shape, offset[..., None])
1087 if is_binary
1088 else lazy(
1089 jnp.full, (binary_indices.size, n), offset[binary_indices, None]
1090 )
1091 if binary_indices is not None
1092 else None
1093 ),
1094 binary_indices=binary_indices,
1095 resid=(
1096 lazy(jnp.zeros, y.shape, storage.resid_dtype)
1097 if is_binary
1098 # resid is created later after y and offset are sharded
1099 else cast(Array, None)
1100 ),
1101 resid_unit=storage.resid_unit,
1102 resid_eff_scale=lazy(jnp.full, kshape, storage.resid_unit),
1103 resid_inexact_integral=lazy(jnp.zeros, kshape),
1104 # only `value` carries the chain axis, so it becomes the lazy leaf;
1105 # the prior params `nu`/`rate` are shared across chains
1106 error_cov_inv=replace(
1107 error_cov_inv, value=_lazy_from_array(error_cov_inv.value)
1108 ),
1109 # `error_scale` goes straight to its field; `missing` is parked in
1110 # the `inv_sdev_scale` slot so it gets sharded with everything
1111 # else. `compute_scale_related_attrs` derives `prec_scale` and
1112 # `inv_sdev_scale` post-shard.
1113 error_scale=error_scale,
1114 prec_scale=None,
1115 inv_sdev_scale=missing,
1116 # invalid placeholders; the true values are set post-shard by
1117 # `compute_scale_related_attrs`
1118 inv_sdev_unit=cast(Array, None),
1119 n_non_missing=cast(Array, None),
1120 sum_diag_prec_scale=cast(Array, None),
1121 forest=Forest(
1122 leaf_tree=lazy(
1123 jnp.zeros, (num_trees, *kshape, tree_size), storage.leaf_dtype
1124 ),
1125 leaf_unit=storage.leaf_unit,
1126 offset=offset,
1127 var_tree=lazy(
1128 jnp.zeros,
1129 (num_trees, tree_size // 2),
1130 minimal_unsigned_dtype(p - 1),
1131 ),
1132 split_tree=lazy(
1133 jnp.zeros, (num_trees, tree_size // 2), max_split.dtype
1134 ),
1135 affluence_tree=lazy(
1136 initial_affluence_tree,
1137 (num_trees, tree_size // 2),
1138 n,
1139 min_points_per_decision_node,
1140 ),
1141 blocked_vars=get_blocked_vars(filter_splitless_vars, max_split),
1142 max_split=max_split,
1143 grow_prop_count=lazy(jnp.zeros, (), int),
1144 grow_acc_count=lazy(jnp.zeros, (), int),
1145 prune_prop_count=lazy(jnp.zeros, (), int),
1146 prune_acc_count=lazy(jnp.zeros, (), int),
1147 p_nonterminal=p_nonterminal[tree_depths(tree_size)],
1148 p_propose_grow=p_nonterminal[tree_depths(tree_size // 2)],
1149 leaf_indices=lazy(
1150 jnp.ones, (num_trees, n), minimal_unsigned_dtype(tree_size - 1)
1151 ),
1152 # no pending prune: `step` starts by applying the pending
1153 # prunes to `leaf_indices`, which shall be a no-op on init
1154 to_prune=lazy(jnp.zeros, (num_trees,), bool),
1155 move_node=lazy(jnp.zeros, (num_trees,), jnp.int32),
1156 # the counts serve the minimum-points constraints and stand in
1157 # for the precisions when there are no per-datapoint scales
1158 # (`prec_scale` is set iff `error_scale` or `missing` is given)
1159 count_tree=(
1160 lazy(initial_count_tree, (num_trees, tree_size), n)
1161 if min_points_per_decision_node is not None
1162 or min_points_per_leaf is not None
1163 or (error_scale is None and missing is None)
1164 else None
1165 ),
1166 # prec_tree is created later, it needs the sharded prec_scale
1167 prec_tree=None,
1168 min_points_per_decision_node=asarray_or_none(
1169 min_points_per_decision_node
1170 ),
1171 min_points_per_leaf=asarray_or_none(min_points_per_leaf),
1172 log_trans_prior=lazy(jnp.zeros, (num_trees,)) if save_ratios else None,
1173 log_likelihood=lazy(jnp.zeros, (num_trees,)) if save_ratios else None,
1174 # only `value` carries the chain axis, see `error_cov_inv`
1175 leaf_prior_cov_inv=replace(
1176 leaf_prior_cov_inv, value=_lazy_from_array(leaf_prior_cov_inv.value)
1177 ),
1178 log_s=_lazy_from_array(asarray_or_none(log_s)),
1179 theta=_lazy_from_array(asarray_or_none(theta)),
1180 rho=asarray_or_none(rho),
1181 a=asarray_or_none(a),
1182 b=asarray_or_none(b),
1183 ),
1184 config=StepConfig(
1185 steps_done=jnp.int32(0),
1186 sparse_on_at=asarray_or_none(sparse_on_at),
1187 sequential_unroll=sequential_unroll,
1188 augment=augment,
1189 mesh=mesh,
1190 leaf_quantization=asarray_or_none(leaf_quantization),
1191 **red_cfg,
1192 ),
1193 )
1195 # add the chain axis to every chain-marked leaf at the position
1196 # declared by its field metadata
1197 state = add_chains(state, num_chains)
1199 # delete big input arrays such that they can be deleted as soon as they
1200 # are sharded, only those arrays that contain an (n,) sized axis
1201 del X, error_scale, missing, y
1203 # move all arrays to the appropriate device and instantiate lazy arrays
1204 state = shard_state(state)
1206 # replace y at masked positions post-shard (the mask is parked in
1207 # `inv_sdev_scale`), before `resid` is derived from y
1208 if state.inv_sdev_scale is not None:
1209 state = replace(
1210 state, y=sanitize_y(state.y, state.inv_sdev_scale, state.forest.offset)
1211 )
1213 # derive the scale arrays and their constant summaries after sharding
1214 # to do the calculation on the right devices. `state.error_scale`
1215 # already holds the sharded user-supplied scale and
1216 # `state.inv_sdev_scale` holds the parked `missing` mask.
1217 # `compute_scale_related_attrs` does not donate `error_scale`, so it
1218 # stays in place; the derived scales fold in the mask, the raw scale
1219 # does not.
1220 attrs = compute_scale_related_attrs(
1221 state.error_scale, state.inv_sdev_scale, storage.prec_scale_dtype, n
1222 )
1223 state = replace(
1224 state,
1225 prec_scale=attrs.prec_scale,
1226 inv_sdev_scale=attrs.inv_sdev_scale,
1227 inv_sdev_unit=attrs.inv_sdev_unit,
1228 n_non_missing=attrs.n_non_missing,
1229 sum_diag_prec_scale=attrs.sum_diag_prec_scale,
1230 )
1232 # calculate initial resid in the continuous outcome case, such that y
1233 # and offset are already sharded if needed
1234 if state.resid is None:
1235 state = set_initial_resid(
1236 state, binary_indices, num_chains, storage.resid_dtype
1237 )
1238 # charge the one-time rounding of the initial residuals into storage
1239 # when it uses a dtype narrower than y (else the cast is exact)
1240 if jnp.finfo(storage.resid_dtype).nmant < jnp.finfo(state.y.dtype).nmant:
1241 state = replace(
1242 state,
1243 resid_inexact_integral=initial_resid_inexact_integral(
1244 state.resid, state.n_non_missing, num_trees
1245 ),
1246 )
1248 # calculate the initial prec_tree from the sharded prec_scale
1249 if state.prec_scale is not None:
1250 state = set_initial_prec_tree(state, num_chains, num_trees, tree_size)
1252 # all the wrong-typed intermediates have now been replaced by their final
1253 # values, so type-checking can resume; make all types strong to avoid
1254 # unwanted recompilations
1255 return remove_weak_types(state)
1258def parse_float_dtype(dtype: DTypeLike) -> jnp.dtype:
1259 """Normalize a storage dtype and check it is floating point."""
1260 dtype = jnp.dtype(dtype)
1261 assert jnp.issubdtype(dtype, jnp.floating)
1262 return dtype
1265def compute_leaf_unit(
1266 leaf_prior_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'],
1267 kshape: tuple[int, ...],
1268) -> Float32[Array, ''] | Float32[Array, ' k']:
1269 """Compute the marginal prior standard deviation of a leaf.
1271 A degenerate prior precision (e.g., infinite, from `bartz.Bart` with
1272 constant ``y``) yields a zero or non-finite scale; fall back to 1 to avoid
1273 nan leaves (an infinite precision pins the leaves to zero anyway).
1274 """
1275 if kshape:
1276 leaf_prior_cov = inv_via_chol_with_gersh(leaf_prior_cov_inv)
1277 leaf_unit = jnp.sqrt(jnp.diagonal(leaf_prior_cov))
1278 else:
1279 leaf_unit = jnp.sqrt(jnp.reciprocal(leaf_prior_cov_inv))
1280 return jnp.where(jnp.isfinite(leaf_unit) & (leaf_unit > 0), leaf_unit, 1.0)
1283def round_to_pow2(
1284 x: Float32[Array, ''] | Float32[Array, ' k'],
1285) -> Float32[Array, ''] | Float32[Array, ' k']:
1286 """Round to the nearest power of two."""
1287 # note: don't use exp2, not exact. `2 ** x` checked exact on cpu & cuda.
1288 return 2 ** jnp.round(jnp.log2(x))
1291def scaled_error_cov_inv(
1292 state: State,
1293) -> Float32[Array, '*chains'] | Float32[Array, '*chains k k']:
1294 """Return the error precision with ``inv_sdev_unit ** 2`` folded in.
1296 `State.prec_scale` and everything summed from it (`Forest.prec_tree`, the
1297 per-leaf precision-scaled residual sums) are stored in units of
1298 ``State.inv_sdev_unit ** 2``, so their products with this scaled precision
1299 are in data units. Returns ``error_cov_inv.value`` as-is when there are no
1300 per-datapoint error scales.
1301 """
1302 value = state.error_cov_inv.value
1303 unit = state.inv_sdev_unit
1304 if state.prec_scale is None:
1305 return value
1306 elif unit.ndim:
1307 return value * (unit[:, None] * unit[None, :])
1308 else:
1309 return value * jnp.square(unit)
1312@dataclass(frozen=True)
1313class StorageParams:
1314 """Storage dtypes and units for the leaves and residuals."""
1316 leaf_dtype: jnp.dtype
1317 prec_scale_dtype: jnp.dtype
1318 resid_dtype: jnp.dtype
1319 leaf_unit: Float32[Array, ''] | Float32[Array, ' k']
1320 resid_unit: Float32[Array, ''] | Float32[Array, ' k']
1323def determine_storage_params(
1324 *,
1325 leaf_dtype: DTypeLike,
1326 prec_scale_dtype: DTypeLike,
1327 resid_dtype: DTypeLike,
1328 leaf_prior_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'],
1329 kshape: tuple[int, ...],
1330 num_trees: int,
1331) -> StorageParams:
1332 """Normalize the storage dtypes and compute the leaf and residual units."""
1333 leaf_unit = compute_leaf_unit(leaf_prior_cov_inv, kshape)
1334 return StorageParams(
1335 leaf_dtype=parse_float_dtype(leaf_dtype),
1336 prec_scale_dtype=parse_float_dtype(prec_scale_dtype),
1337 resid_dtype=parse_float_dtype(resid_dtype),
1338 leaf_unit=round_to_pow2(leaf_unit),
1339 resid_unit=round_to_pow2(leaf_unit * num_trees**0.5),
1340 )
1343def set_initial_resid(
1344 state: 'State',
1345 binary_indices: Int32[Array, ' kb'] | None,
1346 num_chains: int | None,
1347 resid_dtype: jnp.dtype,
1348) -> 'State':
1349 """Build the continuous-outcome `resid` and shard it.
1351 Called post-shard so the captured ``state.y`` and
1352 ``state.forest.offset`` are already on the target devices. Sharding axes are
1353 read via `chain_vmap_axes` / `data_vmap_axes` on a shape preview where the
1354 new `resid` leaf has the chain-extended ``ndim`` (inflated by a placeholder
1355 when `num_chains` is not `None`).
1356 """
1357 inner = _LazyArray(
1358 initial_resid,
1359 state.y.shape,
1360 state.y,
1361 state.forest.offset,
1362 binary_indices,
1363 state.resid_unit,
1364 resid_dtype,
1365 )
1366 preview_resid = add_dummy_axis(inner) if num_chains is not None else inner
1367 preview = replace(state, resid=preview_resid)
1368 chain_axis = chain_vmap_axes(preview).resid
1369 data_axis = data_vmap_axes(preview).resid
1370 resid = _wrap_chain(inner, chain_axis, num_chains)
1371 resid = shard_leaf(resid, chain_axis, data_axis, state.config.mesh)
1372 return replace(state, resid=resid)
1375@jit
1376def initial_resid_inexact_integral(
1377 resid: Float[Array, '*chains n'] | Float[Array, '*chains k n'],
1378 n_non_missing: Int32[Array, ''] | Int32[Array, ' k'],
1379 num_trees: int,
1380) -> Float32[Array, '*chains'] | Float32[Array, '*chains k']:
1381 """Seed value for `State.resid_inexact_integral` from the initial rounding.
1383 Casting the initial residuals to a storage dtype narrower than `y` rounds
1384 them once, and the running updates never fix this offset (they either round
1385 again, which the per-step accounting covers, or preserve it exactly), so it
1386 is charged upfront as one tree-update's worth of rounding.
1387 """
1388 # masked residuals are set to 0 at this stage, so they drop out of the sum
1389 ms = jnp.einsum('...n,...n->...', resid, resid, preferred_element_type=jnp.float32)
1390 ms /= jnp.maximum(n_non_missing, 1)
1391 return ms / num_trees
1394def initial_resid( # noqa: PLR0917, `_LazyArray` passes arguments by position
1395 shape: tuple[int, ...],
1396 y: Float32[Array, ' n'] | Float32[Array, 'k n'],
1397 offset: Float32[Array, ''] | Float32[Array, ' k'],
1398 binary_indices: Int32[Array, ' kb'] | None,
1399 resid_unit: Float32[Array, ''] | Float32[Array, ' k'],
1400 resid_dtype: jnp.dtype,
1401) -> Float[Array, ' n'] | Float[Array, 'k n']:
1402 """Calculate the initial value for `State.resid` in the continuous outcome case.
1404 The residual is stored in units of `resid_unit` and dtype `resid_dtype`. In
1405 the mixed binary-continuous case, binary rows are zeroed out (their residual
1406 starts at ``z - trees - offset = 0``).
1407 """
1408 resid = jnp.broadcast_to(y - offset[..., None], shape)
1409 if binary_indices is not None:
1410 resid = resid.at[..., binary_indices, :].set(0.0)
1411 return (resid / resid_unit[..., None]).astype(resid_dtype)
1414def initial_affluence_tree(
1415 shape: tuple[int, ...], n: int, min_points_per_decision_node: int | None
1416) -> Shaped[Array, '...']:
1417 """Create the initial value of `Forest.affluence_tree`."""
1418 return (
1419 jnp.zeros(shape, bool)
1420 .at[..., 1]
1421 .set(
1422 True
1423 if min_points_per_decision_node is None
1424 else n >= min_points_per_decision_node
1425 )
1426 )
1429def initial_count_tree(shape: tuple[int, ...], n: int) -> Shaped[Array, '...']:
1430 """Create the initial value of `Forest.count_tree`: all datapoints in the root."""
1431 return jnp.zeros(shape, jnp.uint32).at[..., 1].set(n)
1434def set_initial_prec_tree(
1435 state: State, num_chains: int | None, num_trees: int, tree_size: int
1436) -> State:
1437 """Build the cached per-leaf precision for root-only trees and shard it.
1439 Called post-shard so the captured ``state.prec_scale`` is already on the
1440 target devices; mirrors `set_initial_resid`.
1441 """
1442 assert state.prec_scale is not None
1443 shape = (num_trees, *state.prec_scale.shape[:-1], tree_size)
1444 inner = _LazyArray(initial_prec_tree, shape, state.prec_scale)
1445 preview_tree = add_dummy_axis(inner) if num_chains is not None else inner
1446 preview = replace(state, forest=replace(state.forest, prec_tree=preview_tree))
1447 chain_axis = chain_vmap_axes(preview).forest.prec_tree
1448 prec_tree = _wrap_chain(inner, chain_axis, num_chains)
1449 prec_tree = shard_leaf(prec_tree, chain_axis, None, state.config.mesh)
1450 return replace(state, forest=replace(state.forest, prec_tree=prec_tree))
1453def initial_prec_tree(
1454 shape: tuple[int, ...], prec_scale: Float[Array, ' n'] | Float[Array, 'k k n']
1455) -> Float32[Array, 'num_trees tree_size'] | Float32[Array, 'num_trees k k tree_size']:
1456 """Create the initial value of `Forest.prec_tree`: all datapoints in the root."""
1457 return (
1458 jnp.zeros(shape, jnp.float32)
1459 .at[..., 1]
1460 .set(prec_scale.sum(axis=-1, dtype=jnp.float32))
1461 )
1464@jit(donate_argnums=(0,))
1465def sanitize_y(
1466 y: Float32[Array, ' n'] | Float32[Array, 'k n'],
1467 missing: Bool[Array, ' n'] | Bool[Array, 'k n'],
1468 offset: Float32[Array, ''] | Float32[Array, ' k'],
1469) -> Float32[Array, ' n'] | Float32[Array, 'k n']:
1470 """Replace `y` with `offset` at masked positions.
1472 The MCMC ignores masked datapoints through their zeroed precision, but the
1473 values parked in `y` still enter `resid`; garbage values of large magnitude
1474 would degrade the accuracy of quantities derived from it, like the ``y -
1475 resid`` train predictions. Donates `y` to overwrite it in place.
1476 """
1477 return jnp.where(missing, offset[..., None], y)
1480class ScaleRelatedAttrs(NamedTuple):
1481 """Output of `compute_scale_related_attrs`."""
1483 inv_sdev_scale: Float[Array, ' n'] | Float[Array, 'k n'] | None
1484 prec_scale: Float[Array, ' n'] | Float[Array, 'k k n'] | None
1485 inv_sdev_unit: Float32[Array, ''] | Float32[Array, ' k']
1486 n_non_missing: Int32[Array, ''] | Int32[Array, ' k']
1487 sum_diag_prec_scale: Float32[Array, ''] | Float32[Array, ' k']
1490@jit(donate_argnums=(1,), static_argnums=(2, 3))
1491def compute_scale_related_attrs(
1492 error_scale: Float32[Array, ' n'] | Float32[Array, 'k n'] | None,
1493 missing: Bool[Array, ' n'] | Bool[Array, 'k n'] | None,
1494 prec_scale_dtype: jnp.dtype,
1495 n: int,
1496) -> ScaleRelatedAttrs:
1497 """Compute all the fixed quantities derived from `error_scale` and `missing`."""
1498 if error_scale is None and missing is None:
1499 n_non_missing = jnp.full((), n)
1500 return ScaleRelatedAttrs(
1501 inv_sdev_scale=None,
1502 prec_scale=None,
1503 inv_sdev_unit=jnp.float32(1.0),
1504 n_non_missing=n_non_missing,
1505 sum_diag_prec_scale=n_non_missing.astype(jnp.float32),
1506 )
1508 # compute inv_sdev_scale
1509 if error_scale is None:
1510 inv_sdev_scale = jnp.array(1.0) # becomes a vector with the `where` below
1511 else:
1512 inv_sdev_scale = jnp.reciprocal(error_scale)
1513 if missing is not None:
1514 inv_sdev_scale = jnp.where(missing, 0.0, inv_sdev_scale)
1516 # count non-missing points, and the equivalent error-scaled quantity
1517 n_non_missing = jnp.sum(inv_sdev_scale != 0, axis=-1)
1518 sum_diag_prec_scale = jnp.einsum('...n,...n->...', inv_sdev_scale, inv_sdev_scale)
1520 # compute inv_sdev_unit and factor it out of inv_sdev_scale
1521 inv_sdev_unit = round_to_pow2(
1522 jnp.sqrt(sum_diag_prec_scale / jnp.maximum(n_non_missing, 1))
1523 )
1524 inv_sdev_unit = jnp.where(inv_sdev_unit, inv_sdev_unit, 1.0)
1525 inv_sdev_scale = inv_sdev_scale / inv_sdev_unit[..., None]
1527 # compute prec_scale
1528 if inv_sdev_scale.ndim == 1:
1529 prec_scale = jnp.square(inv_sdev_scale)
1530 else:
1531 prec_scale = jnp.einsum('an,bn->abn', inv_sdev_scale, inv_sdev_scale)
1533 return ScaleRelatedAttrs(
1534 inv_sdev_scale=inv_sdev_scale.astype(prec_scale_dtype),
1535 prec_scale=prec_scale.astype(prec_scale_dtype),
1536 inv_sdev_unit=inv_sdev_unit,
1537 n_non_missing=n_non_missing,
1538 sum_diag_prec_scale=sum_diag_prec_scale,
1539 )
1542def get_blocked_vars(
1543 filter_splitless_vars: int, max_split: UInt[Array, ' p']
1544) -> None | UInt[Array, ' q']:
1545 """Initialize the `blocked_vars` field."""
1546 if filter_splitless_vars:
1547 (p,) = max_split.shape
1548 (blocked_vars,) = jnp.nonzero(
1549 max_split == 0, size=filter_splitless_vars, fill_value=p
1550 )
1551 return blocked_vars.astype(minimal_unsigned_dtype(p))
1552 # see `fully_used_variables` for the type cast
1553 else:
1554 return None
1557def add_chains(state: 'State', num_chains: int | None) -> 'State':
1558 """Extend chain-marked `_LazyArray` leaves to include a chain axis of size `num_chains`.
1560 Walks `state`, asks `chain_vmap_axes` where each leaf's chain axis lives,
1561 and wraps the carried `_LazyArray` so its factory creates the core array
1562 and then broadcasts a chain axis in at that position. To make
1563 `chain_vmap_axes` normalize against the chain-extended ``ndim``, the
1564 lookup is done on a shape preview built via `add_dummy_axis`. No-op when
1565 `num_chains` is `None`.
1567 Chain-marked leaves are required to be `_LazyArray` (or `None`); eager
1568 arrays at chain-marked positions are rejected so that all chain insertion
1569 happens at concretization time inside `shard_state`.
1570 """
1571 if num_chains is None:
1572 return state
1573 preview = add_dummy_axis(state)
1574 chain_axes = chain_vmap_axes(preview)
1576 def wrap(leaf: object, chain_axis: int | None) -> object:
1577 if chain_axis is None or leaf is None:
1578 return leaf
1579 assert isinstance(leaf, _LazyArray), (
1580 f'expected _LazyArray for chain-marked leaf, got {type(leaf).__name__}'
1581 )
1582 return _wrap_chain(leaf, chain_axis, num_chains)
1584 return tree.map(wrap, state, chain_axes, is_leaf=_is_lazy_or_none)
1587def parse_mesh(
1588 num_chains: int | None, mesh: Mesh | dict[str, int] | None
1589) -> Mesh | None:
1590 """Parse the `mesh` argument."""
1591 if mesh is None:
1592 return None
1594 # convert dict format to actual mesh
1595 if not isinstance(mesh, Mesh):
1596 assert set(mesh).issubset({'chains', 'data'})
1597 mesh = make_mesh(
1598 tuple(mesh.values()), tuple(mesh), axis_types=(AxisType.Auto,) * len(mesh)
1599 )
1601 # the chains mesh axis must be consistent with the number of chains
1602 if 'chains' in mesh.axis_names:
1603 if num_chains is None:
1604 msg = "mesh has a 'chains' axis but num_chains is None (scalar, no chain axis)"
1605 raise ValueError(msg)
1606 chains_axis = get_axis_size(mesh, 'chains')
1607 if num_chains % chains_axis:
1608 msg = (
1609 f"mesh 'chains' axis of size {chains_axis} does not divide "
1610 f'num_chains={num_chains}'
1611 )
1612 raise ValueError(msg)
1614 # check the axes we use are in auto mode
1615 assert 'chains' not in mesh.axis_names or 'chains' in mesh.auto_axes
1616 assert 'data' not in mesh.axis_names or 'data' in mesh.auto_axes
1618 return mesh
1621@partial(filter_jit, donate='all')
1622# jit and donate because otherwise type conversion would create copies
1623def remove_weak_types(x: PyTree[Array, 'T']) -> PyTree[Array, 'T']:
1624 """Make all types strong.
1626 This is to avoid recompilation in `run_mcmc` or `step`.
1627 """
1629 def remove_weak(x: T) -> T:
1630 if isinstance(x, Array) and x.weak_type:
1631 return cast(T, x.astype(x.dtype))
1632 else:
1633 return x
1635 return tree.map(remove_weak, x)
1638def shard_state(state: State) -> State:
1639 """Place all arrays on the appropriate devices, and instantiate lazily defined arrays."""
1640 mesh = state.config.mesh
1641 shard_leaf_mesh = partial(shard_leaf, mesh=mesh)
1642 return tree.map(
1643 shard_leaf_mesh,
1644 state,
1645 chain_vmap_axes(state),
1646 data_vmap_axes(state),
1647 is_leaf=lambda x: x is None or isinstance(x, _LazyArray),
1648 )
1651def leaf_partition_spec(
1652 ndim: int, chain_axis: int | None, data_axis: int | None, mesh: Mesh
1653) -> PartitionSpec:
1654 """Build a `PartitionSpec` for a leaf with the given chain/data axes."""
1655 spec = [None] * ndim
1656 if chain_axis is not None and 'chains' in mesh.axis_names:
1657 spec[chain_axis] = 'chains'
1658 if data_axis is not None and 'data' in mesh.axis_names:
1659 spec[data_axis] = 'data'
1661 # remove trailing Nones to be consistent with jax's output, it's useful
1662 # for comparing shardings during debugging
1663 while spec and spec[-1] is None:
1664 spec.pop()
1666 return PartitionSpec(*spec)
1669def shard_leaf(
1670 x: Shaped[Array, '*shape'] | None | Shaped[_LazyArray, '*shape'],
1671 chain_axis: int | None,
1672 data_axis: int | None,
1673 mesh: Mesh | None,
1674) -> Shaped[Array, '*shape'] | None:
1675 """Create `x` if it's lazy and shard it."""
1676 if x is None:
1677 return None
1679 if mesh is None:
1680 sharding = None
1681 else:
1682 spec = leaf_partition_spec(x.ndim, chain_axis, data_axis, mesh)
1683 sharding = NamedSharding(mesh, spec)
1685 if isinstance(x, _LazyArray):
1686 x = concretize_lazy_array(x, sharding)
1687 elif sharding is not None:
1688 x = device_put(x, sharding, donate=True)
1690 return x
1693@filter_jit
1694# jit such that in recent jax versions the shards are created on the right
1695# devices immediately instead of being created on the wrong device and then
1696# copied
1697def concretize_lazy_array(
1698 x: Shaped[_LazyArray, '*shape'], sharding: NamedSharding | None
1699) -> Shaped[Array, '*shape']:
1700 """Create an array from an abstract spec on the appropriate devices."""
1701 x = x()
1702 if sharding is not None:
1703 x = lax.with_sharding_constraint(x, sharding)
1704 return x
1707def all_none_or_not_none(*args: object) -> bool:
1708 is_none = [x is None for x in args]
1709 return all(is_none) or not any(is_none)
1712def asarray_or_none(x: object) -> Shaped[Array, '...'] | None:
1713 if x is None:
1714 return None
1715 return jnp.asarray(x)
1718class ReductionConfigs(TypedDict):
1719 """Fields of `StepConfig` related to reductions."""
1721 resid_reduction_config: ReductionConfig
1722 count_reduction_config: ReductionConfig
1723 prec_reduction_config: ReductionConfig
1724 prec_count_num_trees: int | None
1727def parse_reduction_configs(
1728 *,
1729 resid_reduction_config: ReductionConfig,
1730 count_reduction_config: ReductionConfig,
1731 prec_reduction_config: ReductionConfig,
1732 prec_count_num_trees: int | None | Literal['auto'],
1733 y: Float32[Array, ' n'] | Float32[Array, ' k n'] | Bool[Array, ' n'],
1734 num_trees: int,
1735 num_chains: int | None,
1736 mesh: Mesh | None,
1737) -> ReductionConfigs:
1738 """Determine settings for indexed reduces."""
1739 n = y.shape[-1]
1740 n //= get_axis_size(mesh, 'data') # per-device datapoints
1742 # chains are vmapped together on each device, so they share the per-step
1743 # memory of the per-tree reduction
1744 chains_per_device = (num_chains or 1) // get_axis_size(mesh, 'chains')
1746 # the reduction configs carry their own datapoint-batch settings (resolved
1747 # per-platform at run time when 'auto', see `ReductionConfig`), so they are
1748 # stored verbatim; only `prec_count_num_trees`, which does not depend on the
1749 # platform, is resolved here
1750 return dict(
1751 resid_reduction_config=resid_reduction_config,
1752 count_reduction_config=count_reduction_config,
1753 prec_reduction_config=prec_reduction_config,
1754 prec_count_num_trees=parse_prec_count_num_trees(
1755 prec_count_num_trees, num_trees, n * chains_per_device
1756 ),
1757 )
1760def parse_prec_count_num_trees(
1761 prec_count_num_trees: int | None | Literal['auto'], num_trees: int, n: int
1762) -> int | None:
1763 """Return the number of trees to process at a time or determine it automatically."""
1764 if prec_count_num_trees != 'auto':
1765 return prec_count_num_trees
1766 max_n_by_ntree = 2**27 # about 100M
1767 pcnt = max_n_by_ntree // max(1, n)
1768 pcnt = min(num_trees, pcnt)
1769 pcnt = max(1, pcnt)
1770 pcnt = search_divisor(
1771 pcnt, num_trees, max(1, pcnt // 2), max(1, min(num_trees, pcnt * 2))
1772 )
1773 if pcnt >= num_trees: 1773 ↛ 1775line 1773 didn't jump to line 1775 because the condition on line 1773 was always true
1774 pcnt = None
1775 return pcnt
1778def search_divisor(target_divisor: int, dividend: int, low: int, up: int) -> int:
1779 """Find the divisor closest to `target_divisor` in [low, up] if `target_divisor` is not already.
1781 If there is none, give up and return `target_divisor`.
1782 """
1783 assert target_divisor >= 1
1784 assert 1 <= low <= up <= dividend
1785 if dividend % target_divisor == 0:
1786 return target_divisor
1787 candidates = numpy.arange(low, up + 1)
1788 divisors = candidates[dividend % candidates == 0]
1789 if divisors.size == 0:
1790 return target_divisor
1791 penalty = numpy.abs(divisors - target_divisor)
1792 closest = numpy.argmin(penalty)
1793 return divisors[closest].item()
1796def get_axis_size(mesh: Mesh | None, axis_name: str) -> int:
1797 if mesh is None or axis_name not in mesh.shape:
1798 return 1
1799 else:
1800 return mesh.shape[axis_name]
1803def chol_with_gersh(
1804 mat: Float32[Array, '*batch_shape k k'], absolute_eps: bool = False
1805) -> Float32[Array, '*batch_shape k k']:
1806 """Cholesky with Gershgorin stabilization, supports batching."""
1807 return _chol_with_gersh_impl(mat, absolute_eps)
1810@partial(jnp.vectorize, signature='(k,k)->(k,k)', excluded=(1,))
1811def _chol_with_gersh_impl(
1812 mat: Float32[Array, '*batch_shape k k'], absolute_eps: bool
1813) -> Float32[Array, '*batch_shape k k']:
1814 # standardize to unit diagonal first, so the Gershgorin shift is relative to
1815 # each component's scale instead of a single absolute value set by the
1816 # largest one (which would swamp components with much smaller variance, e.g.
1817 # mixing a heavily-scaled continuous outcome with O(1) binary ones).
1818 # degenerate (non-positive or non-finite, e.g. infinite precision from a
1819 # constant outcome) diagonals fall back to the largest finite scale, which
1820 # keeps a diagonal matrix bit-identical to an unstandardized stabilization
1821 # and leaves an infinite diagonal infinite (pinning that leaf to zero)
1822 diag = jnp.diagonal(mat)
1823 finite_pos = jnp.isfinite(diag) & (diag > 0)
1824 ref = jnp.max(jnp.where(finite_pos, diag, 0.0), initial=0.0)
1825 scale = jnp.sqrt(jnp.where(finite_pos, diag, jnp.where(ref > 0, ref, 1.0)))
1826 mat = mat / (scale[:, None] * scale[None, :])
1827 rho = jnp.max(jnp.sum(jnp.abs(mat), axis=1), initial=0.0)
1828 eps = jnp.finfo(mat.dtype).eps
1829 u = mat.shape[0] * rho * eps
1830 if absolute_eps:
1831 u += eps
1832 mat = mat.at[jnp.diag_indices_from(mat)].add(u)
1833 return scale[:, None] * jnp.linalg.cholesky(mat)
1836def inv_via_chol_with_gersh(
1837 mat: Float32[Array, '*batch_shape k k'],
1838) -> Float32[Array, '*batch_shape k k']:
1839 """Compute matrix inverse via Cholesky with Gershgorin stabilization.
1841 DO NOT USE THIS FUNCTION UNLESS YOU REALLY NEED TO.
1842 """
1843 # mat = L L^T
1844 # mat^-1 = L^-T L^-1 = L^-T I L^-1 = L^-T (L^-T I)^T
1845 # I suspect this to be more accurate than (L^-1 I)^T (L^-1 I)
1846 L = chol_with_gersh(mat)
1847 eye = jnp.broadcast_to(jnp.eye(mat.shape[-1]), mat.shape)
1848 Ltinv = solve_triangular(L, eye, trans='T', lower=True)
1849 return solve_triangular(L, Ltinv.mT, trans='T', lower=True)
1852# The chain decorators below are generic over the state type, so that they
1853# preserve `State` subclasses.
1854# WORKAROUND(python<3.12): replace with PEP 695 type parameters on the
1855# decorators, e.g. `def vmap_chains[S: State](...)`.
1856StateT = TypeVar('StateT', bound=State)
1859def split_key_for_chains(
1860 fun: Callable[[Key[Array, ''] | Key[Array, ' num_chains'], StateT], StateT],
1861) -> Callable[[Key[Array, ''], StateT], StateT]:
1862 """Split a single PRNG key into per-chain keys before calling `fun`.
1864 When the state is multichain, the input key is split into
1865 ``state.num_chains()`` keys. For single-chain states, the key is passed
1866 through unchanged.
1867 """
1869 @wraps(fun)
1870 def wrapped(key: Key[Array, ''], state: StateT) -> StateT:
1871 num_chains = state.num_chains()
1872 if num_chains is None:
1873 return fun(key, state)
1874 keys = random.split(key, num_chains)
1875 return fun(keys, state)
1877 return wrapped
1880def partition_specs(x: PyTree, mesh: Mesh) -> PyTree[PartitionSpec]:
1881 """Per-leaf `PartitionSpec`s derived from chain/data `field` markers.
1883 Each array leaf is sharded over ``'chains'`` along its chain axis and over
1884 ``'data'`` along its data axis, when those axes are marked (see `field`)
1885 and present in `mesh`; all other axes are replicated.
1887 Parameters
1888 ----------
1889 x
1890 A pytree of arrays carrying chain/data `field` markers.
1891 mesh
1892 The device mesh to shard over.
1894 Returns
1895 -------
1896 A pytree matching `x` with a `PartitionSpec` in place of each array leaf.
1897 """
1898 return tree.map(
1899 lambda leaf, ca, da: leaf_partition_spec(leaf.ndim, ca, da, mesh),
1900 x,
1901 chain_vmap_axes(x),
1902 data_vmap_axes(x),
1903 )
1906def shard_map_state(
1907 fun: Callable[[Key[Array, ''] | Key[Array, ' num_chains'], StateT], StateT],
1908) -> Callable[[Key[Array, ''] | Key[Array, ' num_chains'], StateT], StateT]:
1909 """Wrap a ``(keys, state) -> state`` function in a manual `jax.shard_map`.
1911 Uses `state.config.mesh` (static). No-op when the mesh is `None`. The keys
1912 input is sharded across ``'chains'`` when the state is multichain and
1913 ``'chains'`` is in the mesh; otherwise the keys are replicated. State
1914 leaves are sharded according to their `chains`/`data` field metadata. The
1915 output sharding matches the input sharding.
1916 """
1918 @wraps(fun)
1919 def wrapped(
1920 key: Key[Array, ''] | Key[Array, ' num_chains'], state: StateT
1921 ) -> StateT:
1922 mesh = state.config.mesh
1923 if mesh is None:
1924 return fun(key, state)
1926 if state.has_chains and 'chains' in mesh.axis_names:
1927 key_spec = PartitionSpec('chains')
1928 else:
1929 key_spec = PartitionSpec()
1931 state_specs = partition_specs(state, mesh)
1933 mapped = shard_map(
1934 fun,
1935 mesh=mesh,
1936 in_specs=(key_spec, state_specs),
1937 out_specs=state_specs,
1938 **get_shard_map_patch_kwargs(),
1939 )
1940 return mapped(key, state)
1942 return wrapped
1945def vmap_chains(
1946 fun: Callable[[Key[Array, ''], StateT], StateT],
1947) -> Callable[[Key[Array, ' num_chains'] | Key[Array, ''], StateT], StateT]:
1948 """Vmap a ``(key, state) -> state`` function over chain axes.
1950 When the state is multichain, `keys` must have a leading chain axis and
1951 `fun` is vmapped over it together with the chain axes of `state`. For
1952 single-chain states, the function is called unchanged.
1953 """
1955 @wraps(fun)
1956 def wrapped(
1957 keys: Key[Array, ' num_chains'] | Key[Array, ''], state: StateT
1958 ) -> StateT:
1959 if not state.has_chains:
1960 return fun(keys, state)
1961 state_axes = chain_vmap_axes(state)
1962 vmapped_fun = vmap(fun, in_axes=(0, state_axes), out_axes=state_axes)
1963 return vmapped_fun(keys, state)
1965 return wrapped
1968class ShardMapPatchKwargs(TypedDict, total=False):
1969 check_vma: bool
1972def get_shard_map_patch_kwargs() -> ShardMapPatchKwargs:
1973 # bug: jax 0.8.1-0.8.2: vmap(shard_map(psum)), jax#34249; the
1974 # jax_disable_vmap_shmap_error config did not work.
1976 # WORKAROUND(jax<=0.8.2): remove this whole function when jax > 0.8.2
1977 buggy = ('0.8.1', '0.8.2')
1978 if jax.__version__ in buggy: 1978 ↛ 1979line 1978 didn't jump to line 1979 because the condition on line 1978 was never true
1979 return {'check_vma': False}
1980 return {}