bartz 0.7.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

bartz/debug.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/debug.py
 #
-# Copyright (c) 2024-2025, Giacomo Petrillo
+# Copyright (c) 2024-2026, The Bartz Contributors
 #
 # This file is part of bartz.
 #
@@ -22,13 +22,14 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-"""Debugging utilities. The entry point is the class `debug_gbart`."""
+"""Debugging utilities. The main functionality is the class `debug_mc_gbart`."""
 
 from collections.abc import Callable
 from dataclasses import replace
 from functools import partial
 from math import ceil, log2
 from re import fullmatch
+from typing import Literal
 
 import numpy
 from equinox import Module, field
@@ -37,20 +38,22 @@ from jax import numpy as jnp
 from jax.tree_util import tree_map
 from jaxtyping import Array, Bool, Float32, Int32, Integer, Key, UInt
 
-from bartz.BART import FloatLike, gbart
+from bartz.BART import gbart, mc_gbart
+from bartz.BART._gbart import FloatLike
 from bartz.grove import (
     TreeHeaps,
     evaluate_forest,
     is_actual_leaf,
     is_leaves_parent,
-    traverse_tree,
+    normalize_axis_tuple,
+    traverse_forest,
     tree_depth,
     tree_depths,
 )
-from bartz.jaxext import minimal_unsigned_dtype, vmap_nodoc
+from bartz.jaxext import autobatch, minimal_unsigned_dtype, vmap_nodoc
 from bartz.jaxext import split as split_key
 from bartz.mcmcloop import TreesTrace
-from bartz.mcmcstep import randint_masked
+from bartz.mcmcstep._moves import randint_masked
 
 
 def format_tree(tree: TreeHeaps, *, print_all: bool = False) -> str:
@@ -154,9 +157,11 @@ def tree_actual_depth(split_tree: UInt[Array, ' 2**(d-1)']) -> Int32[Array, '']:
     return jnp.max(depth)
 
 
+@jit
+@partial(jnp.vectorize, signature='(nt,hts)->(d)')
 def forest_depth_distr(
-    split_tree: UInt[Array, 'num_trees 2**(d-1)'],
-) -> Int32[Array, ' d']:
+    split_tree: UInt[Array, '*batch_shape num_trees 2**(d-1)'],
+) -> Int32[Array, '*batch_shape d']:
     """Histogram the depths of a set of trees.
 
     Parameters
@@ -173,195 +178,102 @@ def forest_depth_distr(
     return jnp.bincount(depths, length=depth)
 
 
-@jit
-def trace_depth_distr(
-    split_tree: UInt[Array, 'trace_length num_trees 2**(d-1)'],
-) -> Int32[Array, 'trace_length d']:
-    """Histogram the depths of a sequence of sets of trees.
-
-    Parameters
-    ----------
-    split_tree
-        The cutpoints of the decision rules of the trees.
-
-    Returns
-    -------
-    A matrix where element (t,i) counts how many trees have depth i in set t.
-    """
-    return vmap(forest_depth_distr)(split_tree)
-
-
-def points_per_decision_node_distr(
-    var_tree: UInt[Array, ' 2**(d-1)'],
-    split_tree: UInt[Array, ' 2**(d-1)'],
+@partial(jit, static_argnames=('node_type', 'sum_batch_axis'))
+def points_per_node_distr(
     X: UInt[Array, 'p n'],
-) -> Int32[Array, ' n+1']:
-    """Histogram points-per-node counts.
-
-    Count how many parent-of-leaf nodes in a tree select each possible amount
-    of points.
-
-    Parameters
-    ----------
-    var_tree
-        The variables of the decision rules.
-    split_tree
-        The cutpoints of the decision rules.
-    X
-        The set of points to count.
-
-    Returns
-    -------
-    A vector where the i-th element counts how many next-to-leaf nodes have i points.
-    """
-    traverse_tree_X = vmap(traverse_tree, in_axes=(1, None, None))
-    indices = traverse_tree_X(X, var_tree, split_tree)
-    indices >>= 1
-    count_tree = jnp.zeros(split_tree.size, int).at[indices].add(1).at[0].set(0)
-    is_parent = is_leaves_parent(split_tree)
-    return jnp.zeros(X.shape[1] + 1, int).at[count_tree].add(is_parent)
-
-
-def forest_points_per_decision_node_distr(
-    trees: TreeHeaps, X: UInt[Array, 'p n']
-) -> Int32[Array, ' n+1']:
-    """Histogram points-per-node counts for a set of trees.
-
-    Count how many parent-of-leaf nodes in a set of trees select each possible
-    amount of points.
-
-    Parameters
-    ----------
-    trees
-        The set of trees. The variables must have broadcast shape (num_trees,).
-    X
-        The set of points to count.
-
-    Returns
-    -------
-    A vector where the i-th element counts how many next-to-leaf nodes have i points.
-    """
-    distr = jnp.zeros(X.shape[1] + 1, int)
+    var_tree: UInt[Array, '*batch_shape 2**(d-1)'],
+    split_tree: UInt[Array, '*batch_shape 2**(d-1)'],
+    node_type: Literal['leaf', 'leaf-parent'],
+    *,
+    sum_batch_axis: int | tuple[int, ...] = (),
+) -> Int32[Array, '*reduced_batch_shape n+1']:
+    """Histogram points-per-node counts in a set of trees.
 
-    def loop(distr, heaps: tuple[Array, Array]):
-        return distr + points_per_decision_node_distr(*heaps, X), None
-
-    distr, _ = lax.scan(loop, distr, (trees.var_tree, trees.split_tree))
-    return distr
-
-
-@jit
-def trace_points_per_decision_node_distr(
-    trace: TreeHeaps, X: UInt[Array, 'p n']
-) -> Int32[Array, 'trace_length n+1']:
-    """Separately histogram points-per-node counts over a sequence of sets of trees.
-
-    For each set of trees, count how many parent-of-leaf nodes select each
-    possible amount of points.
+    Count how many nodes in a tree select each possible amount of points,
+    over a certain subset of nodes.
 
     Parameters
     ----------
-    trace
-        The sequence of sets of trees. The variables must have broadcast shape
-        (trace_length, num_trees).
     X
         The set of points to count.
-
-    Returns
-    -------
-    A matrix where element (t,i) counts how many next-to-leaf nodes have i points in set t.
-    """
-
-    def loop(_, trace):
-        return None, forest_points_per_decision_node_distr(trace, X)
-
-    _, distr = lax.scan(loop, None, trace)
-    return distr
-
-
-def points_per_leaf_distr(
-    var_tree: UInt[Array, ' 2**(d-1)'],
-    split_tree: UInt[Array, ' 2**(d-1)'],
-    X: UInt[Array, 'p n'],
-) -> Int32[Array, ' n+1']:
-    """Histogram points-per-leaf counts in a tree.
-
-    Count how many leaves in a tree select each possible amount of points.
-
-    Parameters
-    ----------
     var_tree
         The variables of the decision rules.
     split_tree
         The cutpoints of the decision rules.
-    X
-        The set of points to count.
-
-    Returns
-    -------
-    A vector where the i-th element counts how many leaves have i points.
-    """
-    traverse_tree_X = vmap(traverse_tree, in_axes=(1, None, None))
-    indices = traverse_tree_X(X, var_tree, split_tree)
-    count_tree = jnp.zeros(2 * split_tree.size, int).at[indices].add(1)
-    is_leaf = is_actual_leaf(split_tree, add_bottom_level=True)
-    return jnp.zeros(X.shape[1] + 1, int).at[count_tree].add(is_leaf)
-
-
-def forest_points_per_leaf_distr(
-    trees: TreeHeaps, X: UInt[Array, 'p n']
-) -> Int32[Array, ' n+1']:
-    """Histogram points-per-leaf counts over a set of trees.
-
-    Count how many leaves in a set of trees select each possible amount of points.
-
-    Parameters
-    ----------
-    trees
-        The set of trees. The variables must have broadcast shape (num_trees,).
-    X
-        The set of points to count.
+    node_type
+        The type of nodes to consider. Can be:
+
+        'leaf'
+            Count only leaf nodes.
+        'leaf-parent'
+            Count only parent-of-leaf nodes.
+    sum_batch_axis
+        Aggregate the histogram over these batch axes, counting how many nodes
+        have each possible amount of points over subsets of trees instead of
+        in each tree separately.
 
     Returns
     -------
-    A vector where the i-th element counts how many leaves have i points.
+    A vector where the i-th element counts how many nodes have i points.
     """
-    distr = jnp.zeros(X.shape[1] + 1, int)
-
-    def loop(distr, heaps: tuple[Array, Array]):
-        return distr + points_per_leaf_distr(*heaps, X), None
-
-    distr, _ = lax.scan(loop, distr, (trees.var_tree, trees.split_tree))
-    return distr
-
-
-@jit
-def trace_points_per_leaf_distr(
-    trace: TreeHeaps, X: UInt[Array, 'p n']
-) -> Int32[Array, 'trace_length n+1']:
-    """Separately histogram points-per-leaf counts over a sequence of sets of trees.
-
-    For each set of trees, count how many leaves select each possible amount of
-    points.
-
-    Parameters
-    ----------
-    trace
-        The sequence of sets of trees. The variables must have broadcast shape
-        (trace_length, num_trees).
-    X
-        The set of points to count.
-
-    Returns
-    -------
-    A matrix where element (t,i) counts how many leaves have i points in set t.
-    """
-
-    def loop(_, trace):
-        return None, forest_points_per_leaf_distr(trace, X)
+    batch_ndim = var_tree.ndim - 1
+    axes = normalize_axis_tuple(sum_batch_axis, batch_ndim)
+
+    def func(
+        var_tree: UInt[Array, '*batch_shape 2**(d-1)'],
+        split_tree: UInt[Array, '*batch_shape 2**(d-1)'],
+    ) -> Int32[Array, '*reduced_batch_shape n+1']:
+        indices: UInt[Array, '*batch_shape n']
+        indices = traverse_forest(X, var_tree, split_tree)
+
+        @partial(jnp.vectorize, signature='(hts),(n)->(ts_or_hts),(ts_or_hts)')
+        def count_points(
+            split_tree: UInt[Array, '*batch_shape 2**(d-1)'],
+            indices: UInt[Array, '*batch_shape n'],
+        ) -> (
+            tuple[UInt[Array, '*batch_shape 2**d'], Bool[Array, '*batch_shape 2**d']]
+            | tuple[
+                UInt[Array, '*batch_shape 2**(d-1)'],
+                Bool[Array, '*batch_shape 2**(d-1)'],
+            ]
+        ):
+            if node_type == 'leaf-parent':
+                indices >>= 1
+                predicate = is_leaves_parent(split_tree)
+            elif node_type == 'leaf':
+                predicate = is_actual_leaf(split_tree, add_bottom_level=True)
+            else:
+                raise ValueError(node_type)
+            count_tree = jnp.zeros(predicate.size, int).at[indices].add(1).at[0].set(0)
+            return count_tree, predicate
+
+        count_tree, predicate = count_points(split_tree, indices)
+
+        def count_nodes(
+            count_tree: UInt[Array, '*summed_batch_axes half_tree_size'],
+            predicate: Bool[Array, '*summed_batch_axes half_tree_size'],
+        ) -> Int32[Array, ' n+1']:
+            return jnp.zeros(X.shape[1] + 1, int).at[count_tree].add(predicate)
+
+        # vmap count_nodes over non-batched dims
+        for i in reversed(range(batch_ndim)):
+            neg_i = i - var_tree.ndim
+            if i not in axes:
+                count_nodes = vmap(count_nodes, in_axes=neg_i)
+
+        return count_nodes(count_tree, predicate)
+
+    # automatically batch over all batch dimensions
+    max_io_nbytes = 2**27  # 128 MiB
+    out_dim_shift = len(axes)
+    for i in reversed(range(batch_ndim)):
+        if i in axes:
+            out_dim_shift -= 1
+        else:
+            func = autobatch(func, max_io_nbytes, i, i - out_dim_shift)
+    assert out_dim_shift == 0
 
-    _, distr = lax.scan(loop, None, trace)
-    return distr
+    return func(var_tree, split_tree)
 
 
 check_functions = []
@@ -400,29 +312,36 @@ def check_types(tree: TreeHeaps, max_split: UInt[Array, ' p']) -> bool:
     return (
         tree.var_tree.dtype == expected_var_dtype
         and tree.split_tree.dtype == expected_split_dtype
+        and jnp.issubdtype(max_split.dtype, jnp.unsignedinteger)
     )
 
 
 @check
-def check_sizes(tree: TreeHeaps, max_split: UInt[Array, ' p']) -> bool:  # noqa: ARG001
+def check_sizes(tree: TreeHeaps, _max_split: UInt[Array, ' p']) -> bool:
     """Check that array sizes are coherent."""
     return tree.leaf_tree.size == 2 * tree.var_tree.size == 2 * tree.split_tree.size
 
 
 @check
-def check_unused_node(tree: TreeHeaps, max_split: UInt[Array, ' p']) -> Bool[Array, '']:  # noqa: ARG001
+def check_unused_node(
+    tree: TreeHeaps, _max_split: UInt[Array, ' p']
+) -> Bool[Array, '']:
     """Check that the unused node slot at index 0 is not dirty."""
     return (tree.var_tree[0] == 0) & (tree.split_tree[0] == 0)
 
 
 @check
-def check_leaf_values(tree: TreeHeaps, max_split: UInt[Array, ' p']) -> Bool[Array, '']:  # noqa: ARG001
+def check_leaf_values(
+    tree: TreeHeaps, _max_split: UInt[Array, ' p']
+) -> Bool[Array, '']:
     """Check that all leaf values are not inf of nan."""
     return jnp.all(jnp.isfinite(tree.leaf_tree))
 
 
 @check
-def check_stray_nodes(tree: TreeHeaps, max_split: UInt[Array, ' p']) -> Bool[Array, '']:  # noqa: ARG001
+def check_stray_nodes(
+    tree: TreeHeaps, _max_split: UInt[Array, ' p']
+) -> Bool[Array, '']:
     """Check if there is any marked-non-leaf node with a marked-leaf parent."""
     index = jnp.arange(
         2 * tree.split_tree.size,
@@ -446,12 +365,12 @@ def check_rule_consistency(
 
     # initial boundaries of decision rules. use extreme integers instead of 0,
     # max_split to avoid checking if there is something out of bounds.
-    small = jnp.iinfo(jnp.int32).min
-    large = jnp.iinfo(jnp.int32).max
-    lower = jnp.full(max_split.size, small, jnp.int32)
-    upper = jnp.full(max_split.size, large, jnp.int32)
-    # specify the type explicitly, otherwise they are weakly types and get
-    # implicitly converted to split.dtype (typically uint8) in the expressions
+    dtype = tree.split_tree.dtype
+    small = jnp.iinfo(dtype).min
+    large = jnp.iinfo(dtype).max
+    lower = jnp.full(max_split.size, small, dtype)
+    upper = jnp.full(max_split.size, large, dtype)
+    # the split must be in (lower[var], upper[var]]
 
     def _check_recursive(node, lower, upper):
         # read decision rule
@@ -464,20 +383,14 @@ def check_rule_consistency(
         upper_var = upper.at[var].get(mode='fill', fill_value=large)
 
         # check rule is in bounds
-        bad = jnp.where(split, (split <= lower_var) | (split >= upper_var), False)
+        bad = jnp.where(split, (split <= lower_var) | (split > upper_var), False)
 
         # recurse
         if node < tree.var_tree.size // 2:
-            bad |= _check_recursive(
-                2 * node,
-                lower,
-                upper.at[jnp.where(split, var, max_split.size)].set(split),
-            )
-            bad |= _check_recursive(
-                2 * node + 1,
-                lower.at[jnp.where(split, var, max_split.size)].set(split),
-                upper,
-            )
+            idx = jnp.where(split, var, max_split.size)
+            bad |= _check_recursive(2 * node, lower, upper.at[idx].set(split - 1))
+            bad |= _check_recursive(2 * node + 1, lower.at[idx].set(split), upper)
+
         return bad
 
     return ~_check_recursive(1, lower, upper)
@@ -557,30 +470,39 @@ def describe_error(error: int | Integer[Array, '']) -> list[str]:
 
 
 @jit
-@partial(vmap_nodoc, in_axes=(0, None))
 def check_trace(
     trace: TreeHeaps, max_split: UInt[Array, ' p']
-) -> UInt[Array, 'trace_length num_trees']:
-    """Check the validity of a sequence of sets of trees.
+) -> UInt[Array, '*batch_shape']:
+    """Check the validity of a set of trees.
 
     Use `describe_error` to parse the error codes returned by this function.
 
     Parameters
     ----------
     trace
-        The sequence of sets of trees to check. The tree arrays must have
-        broadcast shape (trace_length, num_trees). This object can have
-        additional attributes beyond the tree arrays, they are ignored.
+        The set of trees to check. This object can have additional attributes
+        beyond the tree arrays, they are ignored.
     max_split
         The maximum split value for each variable.
 
     Returns
     -------
-    A matrix of error codes for each tree.
+    A tensor of error codes for each tree.
     """
-    trees = TreesTrace.from_dataclass(trace)
-    check_forest = vmap(check_tree, in_axes=(0, None))
-    return check_forest(trees, max_split)
+    # vectorize check_tree over all batch dimensions
+    unpack_check_tree = lambda l, v, s: check_tree(TreesTrace(l, v, s), max_split)
+    is_mv = trace.leaf_tree.ndim > trace.split_tree.ndim
+    signature = '(k,ts),(hts),(hts)->()' if is_mv else '(ts),(hts),(hts)->()'
+    vec_check_tree = jnp.vectorize(unpack_check_tree, signature=signature)
+
+    # automatically batch over all batch dimensions
+    max_io_nbytes = 2**24  # 16 MiB
+    batch_ndim = trace.split_tree.ndim - 1
+    batched_check_tree = vec_check_tree
+    for i in reversed(range(batch_ndim)):
+        batched_check_tree = autobatch(batched_check_tree, max_io_nbytes, i, i)
+
+    return batched_check_tree(trace.leaf_tree, trace.var_tree, trace.split_tree)
 
 
 def _get_next_line(s: str, i: int) -> tuple[str, int]:
@@ -592,24 +514,19 @@ def _get_next_line(s: str, i: int) -> tuple[str, int]:
 
 
 class BARTTraceMeta(Module):
-    """Metadata of R BART tree traces.
-
-    Parameters
-    ----------
-    ndpost
-        The number of posterior draws.
-    ntree
-        The number of trees in the model.
-    numcut
-        The maximum split value for each variable.
-    heap_size
-        The size of the heap required to store the trees.
-    """
+    """Metadata of R BART tree traces."""
 
     ndpost: int = field(static=True)
+    """The number of posterior draws."""
+
     ntree: int = field(static=True)
+    """The number of trees in the model."""
+
     numcut: UInt[Array, ' p']
+    """The maximum split value for each variable."""
+
     heap_size: int = field(static=True)
+    """The size of the heap required to store the trees."""
 
 
 def scan_BART_trees(trees: str) -> BARTTraceMeta:
@@ -801,25 +718,24 @@ class SamplePriorStack(Module):
 
     Each level of the stack represents a recursion into a child node in a
     binary tree of maximum depth `d`.
-
-    Parameters
-    ----------
-    nonterminal
-        Whether the node is valid or the recursion is into unused node slots.
-    lower
-    upper
-        The available cutpoints along ``var`` are in the integer range
-        ``[1 + lower[var], 1 + upper[var])``.
-    var
-    split
-        The variable and cutpoint of a decision node.
     """
 
     nonterminal: Bool[Array, ' d-1']
+    """Whether the node is valid or the recursion is into unused node slots."""
+
     lower: UInt[Array, 'd-1 p']
+    """The available cutpoints along ``var`` are in the integer range
+    ``[1 + lower[var], 1 + upper[var])``."""
+
     upper: UInt[Array, 'd-1 p']
+    """The available cutpoints along ``var`` are in the integer range
+    ``[1 + lower[var], 1 + upper[var])``."""
+
     var: UInt[Array, ' d-1']
+    """The variable of a decision node."""
+
     split: UInt[Array, ' d-1']
+    """The cutpoint of a decision node."""
 
     @classmethod
     def initial(
@@ -850,19 +766,16 @@ class SamplePriorStack(Module):
 
 
 class SamplePriorTrees(Module):
-    """Object holding the trees generated by `sample_prior`.
-
-    Parameters
-    ----------
-    leaf_tree
-    var_tree
-    split_tree
-        The arrays representing the trees, see `bartz.grove`.
-    """
+    """Object holding the trees generated by `sample_prior`."""
 
     leaf_tree: Float32[Array, '* 2**d']
+    """The array representing the trees, see `bartz.grove`."""
+
     var_tree: UInt[Array, '* 2**(d-1)']
+    """The array representing the trees, see `bartz.grove`."""
+
     split_tree: UInt[Array, '* 2**(d-1)']
+    """The array representing the trees, see `bartz.grove`."""
 
     @classmethod
     def initial(
@@ -903,21 +816,16 @@ class SamplePriorTrees(Module):
 
 
 class SamplePriorCarry(Module):
-    """Object holding values carried along the recursion in `sample_prior`.
-
-    Parameters
-    ----------
-    key
-        A jax random key used to sample decision rules.
-    stack
-        The stack used to manage the recursion.
-    trees
-        The output arrays.
-    """
+    """Object holding values carried along the recursion in `sample_prior`."""
 
     key: Key[Array, '']
+    """A jax random key used to sample decision rules."""
+
     stack: SamplePriorStack
+    """The stack used to manage the recursion."""
+
     trees: SamplePriorTrees
+    """The output arrays."""
 
     @classmethod
     def initial(
@@ -958,21 +866,17 @@ class SamplePriorX(Module):
 
     The sequence of nodes to visit is pre-computed recursively once, unrolling
     the recursion schedule.
-
-    Parameters
-    ----------
-    node
-        The heap index of the node to visit.
-    depth
-        The depth of the node.
-    next_depth
-        The depth of the next node to visit, either the left child or the right
-        sibling of the node or of an ancestor.
     """
 
     node: Int32[Array, ' 2**(d-1)-1']
+    """The heap index of the node to visit."""
+
     depth: Int32[Array, ' 2**(d-1)-1']
+    """The depth of the node."""
+
     next_depth: Int32[Array, ' 2**(d-1)-1']
+    """The depth of the next node to visit, either the left child or the right
+    sibling of the node or of an ancestor."""
 
     @classmethod
     def initial(cls, p_nonterminal: Float32[Array, ' d-1']) -> 'SamplePriorX':
@@ -992,7 +896,7 @@ class SamplePriorX(Module):
         assert len(seq) == 2**p_nonterminal.size - 1
         node = [node for node, depth in seq]
         depth = [depth for node, depth in seq]
-        next_depth = depth[1:] + [p_nonterminal.size]
+        next_depth = [*depth[1:], p_nonterminal.size]
         return cls(
             node=jnp.array(node),
             depth=jnp.array(depth),
@@ -1173,18 +1077,18 @@ def sample_prior(
     return tree_map(lambda x: x.reshape(trace_length, num_trees, -1), trees)
 
 
-class debug_gbart(gbart):
-    """A subclass of `gbart` that adds debugging functionality.
+class debug_mc_gbart(mc_gbart):
+    """A subclass of `mc_gbart` that adds debugging functionality.
 
     Parameters
     ----------
     *args
-        Passed to `gbart`.
+        Passed to `mc_gbart`.
     check_trees
         If `True`, check all trees with `check_trace` after running the MCMC,
         and assert that they are all valid. Set to `False` to allow jax tracing.
     **kw
-        Passed to `gbart`.
+        Passed to `mc_gbart`.
     """
 
     def __init__(self, *args, check_trees: bool = True, **kw):
@@ -1194,24 +1098,28 @@ class debug_gbart(gbart):
             bad_count = jnp.count_nonzero(bad)
             assert bad_count == 0
 
-    def show_tree(self, i_sample: int, i_tree: int, print_all: bool = False):
+    def print_tree(
+        self, i_chain: int, i_sample: int, i_tree: int, print_all: bool = False
+    ):
         """Print a single tree in human-readable format.
 
         Parameters
         ----------
+        i_chain
+            The index of the MCMC chain.
         i_sample
-            The index of the posterior sample.
+            The index of the (post-burnin) sample in the chain.
         i_tree
             The index of the tree in the sample.
         print_all
             If `True`, also print the content of unused node slots.
         """
         tree = TreesTrace.from_dataclass(self._main_trace)
-        tree = tree_map(lambda x: x[i_sample, i_tree, :], tree)
+        tree = tree_map(lambda x: x[i_chain, i_sample, i_tree, :], tree)
         s = format_tree(tree, print_all=print_all)
         print(s)  # noqa: T201, this method is intended for debug
 
-    def sigma_harmonic_mean(self, prior: bool = False) -> Float32[Array, '']:
+    def sigma_harmonic_mean(self, prior: bool = False) -> Float32[Array, ' mc_cores']:
         """Return the harmonic mean of the error variance.
 
         Parameters
@@ -1225,33 +1133,36 @@ class debug_gbart(gbart):
         The harmonic mean 1/E[1/sigma^2] in the selected distribution.
         """
         bart = self._mcmc_state
-        assert bart.sigma2_alpha is not None
+        assert bart.error_cov_df is not None
         assert bart.z is None
+        # inverse gamma prior: alpha = df / 2, beta = scale / 2
         if prior:
-            alpha = bart.sigma2_alpha
-            beta = bart.sigma2_beta
+            alpha = bart.error_cov_df / 2
+            beta = bart.error_cov_scale / 2
         else:
-            resid = bart.resid
-            alpha = bart.sigma2_alpha + resid.size / 2
-            norm2 = resid @ resid
-            beta = bart.sigma2_beta + norm2 / 2
-        sigma2 = beta / alpha
-        return jnp.sqrt(sigma2)
-
-    def compare_resid(self) -> tuple[Float32[Array, ' n'], Float32[Array, ' n']]:
+            alpha = bart.error_cov_df / 2 + bart.resid.size / 2
+            norm2 = jnp.einsum('ij,ij->i', bart.resid, bart.resid)
+            beta = bart.error_cov_scale / 2 + norm2 / 2
+        error_cov_inv = alpha / beta
+        return jnp.sqrt(lax.reciprocal(error_cov_inv))
+
+    def compare_resid(
+        self,
+    ) -> tuple[Float32[Array, 'mc_cores n'], Float32[Array, 'mc_cores n']]:
         """Re-compute residuals to compare them with the updated ones.
 
         Returns
         -------
-        resid1 : Float32[Array, 'n']
+        resid1 : Float32[Array, 'mc_cores n']
             The final state of the residuals updated during the MCMC.
-        resid2 : Float32[Array, 'n']
+        resid2 : Float32[Array, 'mc_cores n']
             The residuals computed from the final state of the trees.
         """
         bart = self._mcmc_state
         resid1 = bart.resid
 
-        trees = evaluate_forest(bart.X, bart.forest)
+        forests = TreesTrace.from_dataclass(bart.forest)
+        trees = evaluate_forest(bart.X, forests, sum_batch_axis=-1)
 
         if bart.z is not None:
             ref = bart.z
@@ -1261,14 +1172,16 @@ class debug_gbart(gbart):
 
         return resid1, resid2
 
-    def avg_acc(self) -> tuple[Float32[Array, ''], Float32[Array, '']]:
+    def avg_acc(
+        self,
+    ) -> tuple[Float32[Array, ' mc_cores'], Float32[Array, ' mc_cores']]:
         """Compute the average acceptance rates of tree moves.
 
         Returns
         -------
-        acc_grow : Float32[Array, '']
+        acc_grow : Float32[Array, 'mc_cores']
             The average acceptance rate of grow moves.
-        acc_prune : Float32[Array, '']
+        acc_prune : Float32[Array, 'mc_cores']
             The average acceptance rate of prune moves.
         """
         trace = self._main_trace
@@ -1276,18 +1189,20 @@ class debug_gbart(gbart):
         def acc(prefix):
             acc = getattr(trace, f'{prefix}_acc_count')
             prop = getattr(trace, f'{prefix}_prop_count')
-            return acc.sum() / prop.sum()
+            return acc.sum(axis=1) / prop.sum(axis=1)
 
         return acc('grow'), acc('prune')
 
-    def avg_prop(self) -> tuple[Float32[Array, ''], Float32[Array, '']]:
+    def avg_prop(
+        self,
+    ) -> tuple[Float32[Array, ' mc_cores'], Float32[Array, ' mc_cores']]:
         """Compute the average proposal rate of grow and prune moves.
 
         Returns
         -------
-        prop_grow : Float32[Array, '']
+        prop_grow : Float32[Array, 'mc_cores']
             The fraction of times grow was proposed instead of prune.
-        prop_prune : Float32[Array, '']
+        prop_prune : Float32[Array, 'mc_cores']
             The fraction of times prune was proposed instead of grow.
 
         Notes
@@ -1298,61 +1213,86 @@ class debug_gbart(gbart):
         trace = self._main_trace
 
         def prop(prefix):
-            return getattr(trace, f'{prefix}_prop_count').sum()
+            return getattr(trace, f'{prefix}_prop_count').sum(axis=1)
 
         pgrow = prop('grow')
         pprune = prop('prune')
         total = pgrow + pprune
         return pgrow / total, pprune / total
 
-    def avg_move(self) -> tuple[Float32[Array, ''], Float32[Array, '']]:
+    def avg_move(
+        self,
+    ) -> tuple[Float32[Array, ' mc_cores'], Float32[Array, ' mc_cores']]:
         """Compute the move rate.
 
         Returns
         -------
-        rate_grow : Float32[Array, '']
+        rate_grow : Float32[Array, 'mc_cores']
             The fraction of times a grow move was proposed and accepted.
-        rate_prune : Float32[Array, '']
+        rate_prune : Float32[Array, 'mc_cores']
             The fraction of times a prune move was proposed and accepted.
         """
         agrow, aprune = self.avg_acc()
         pgrow, pprune = self.avg_prop()
         return agrow * pgrow, aprune * pprune
 
-    def depth_distr(self) -> Float32[Array, 'trace_length d']:
+    def depth_distr(self) -> Int32[Array, 'mc_cores ndpost/mc_cores d']:
         """Histogram of tree depths for each state of the trees.
 
         Returns
         -------
         A matrix where each row contains a histogram of tree depths.
         """
-        return trace_depth_distr(self._main_trace.split_tree)
+        out: Int32[Array, '*chains samples d']
+        out = forest_depth_distr(self._main_trace.split_tree)
+        if out.ndim < 3:
+            out = out[None, :, :]
+        return out
+
+    def _points_per_node_distr(
+        self, node_type: str
+    ) -> Int32[Array, 'mc_cores ndpost/mc_cores n+1']:
+        out: Int32[Array, '*chains samples n+1']
+        out = points_per_node_distr(
+            self._mcmc_state.X,
+            self._main_trace.var_tree,
+            self._main_trace.split_tree,
+            node_type,
+            sum_batch_axis=-1,
+        )
+        if out.ndim < 3:
+            out = out[None, :, :]
+        return out
 
-    def points_per_decision_node_distr(self) -> Float32[Array, 'trace_length n+1']:
+    def points_per_decision_node_distr(
+        self,
+    ) -> Int32[Array, 'mc_cores ndpost/mc_cores n+1']:
         """Histogram of number of points belonging to parent-of-leaf nodes.
 
         Returns
         -------
-        A matrix where each row contains a histogram of number of points.
+        For each chain, a matrix where each row contains a histogram of number of points.
         """
-        return trace_points_per_decision_node_distr(
-            self._main_trace, self._mcmc_state.X
-        )
+        return self._points_per_node_distr('leaf-parent')
 
-    def points_per_leaf_distr(self) -> Float32[Array, 'trace_length n+1']:
+    def points_per_leaf_distr(self) -> Int32[Array, 'mc_cores ndpost/mc_cores n+1']:
         """Histogram of number of points belonging to leaves.
 
         Returns
         -------
         A matrix where each row contains a histogram of number of points.
         """
-        return trace_points_per_leaf_distr(self._main_trace, self._mcmc_state.X)
+        return self._points_per_node_distr('leaf')
 
-    def check_trees(self) -> UInt[Array, 'trace_length ntree']:
+    def check_trees(self) -> UInt[Array, 'mc_cores ndpost/mc_cores ntree']:
         """Apply `check_trace` to all the tree draws."""
-        return check_trace(self._main_trace, self._mcmc_state.forest.max_split)
+        out: UInt[Array, '*chains samples num_trees']
+        out = check_trace(self._main_trace, self._mcmc_state.forest.max_split)
+        if out.ndim < 3:
+            out = out[None, :, :]
+        return out
 
-    def tree_goes_bad(self) -> Bool[Array, 'trace_length ntree']:
+    def tree_goes_bad(self) -> Bool[Array, 'mc_cores ndpost/mc_cores ntree']:
         """Find iterations where a tree becomes invalid.
 
         Returns
@@ -1360,5 +1300,20 @@ class debug_gbart(gbart):
         A where (i,j) is `True` if tree j is invalid at iteration i but not i-1.
         """
         bad = self.check_trees().astype(bool)
-        bad_before = jnp.pad(bad[:-1], [(1, 0), (0, 0)])
+        bad_before = jnp.pad(bad[:, :-1, :], [(0, 0), (1, 0), (0, 0)])
         return bad & ~bad_before
+
+
+class debug_gbart(debug_mc_gbart, gbart):
+    """A subclass of `gbart` that adds debugging functionality.
+
+    Parameters
+    ----------
+    *args
+        Passed to `gbart`.
+    check_trees
+        If `True`, check all trees with `check_trace` after running the MCMC,
+        and assert that they are all valid. Set to `False` to allow jax tracing.
+    **kw
+        Passed to `gbart`.
+    """