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

bartz/grove.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/grove.py
 #
-# Copyright (c) 2024-2025, Giacomo Petrillo
+# Copyright (c) 2024-2026, The Bartz Contributors
 #
 # This file is part of bartz.
 #
@@ -28,10 +28,14 @@ import math
 from functools import partial
 from typing import Protocol
 
-import jax
-from jax import jit, lax
+from jax import jit, lax, vmap
 from jax import numpy as jnp
-from jaxtyping import Array, Bool, DTypeLike, Float32, Int32, Real, Shaped, UInt
+from jaxtyping import Array, Bool, DTypeLike, Float32, Int32, Shaped, UInt
+
+try:
+    from numpy.lib.array_utils import normalize_axis_tuple  # numpy 2
+except ImportError:
+    from numpy.core.numeric import normalize_axis_tuple  # numpy 1
 
 from bartz.jaxext import minimal_unsigned_dtype, vmap_nodoc
 
@@ -44,32 +48,33 @@ class TreeHeaps(Protocol):
     (left child) and :math:`2i + 1` (right child). The array element at index 0
     is unused.
 
-    Parameters
-    ----------
-    leaf_tree
-        The values in the leaves of the trees. This array can be dirty, i.e.,
-        unused nodes can have whatever value.
-    var_tree
-        The axes along which the decision nodes operate. This array can be
-        dirty but for the always unused node at index 0 which must be set to 0.
-    split_tree
-        The decision boundaries of the trees. The boundaries are open on the
-        right, i.e., a point belongs to the left child iff x < split. Whether a
-        node is a leaf is indicated by the corresponding 'split' element being
-        0. Unused nodes also have split set to 0. This array can't be dirty.
-
-    Notes
-    -----
     Since the nodes at the bottom can only be leaves and not decision nodes,
     `var_tree` and `split_tree` are half as long as `leaf_tree`.
+
+    Arrays may have additional initial axes to represent multiple trees.
     """
 
-    leaf_tree: Float32[Array, '* 2**d']
-    var_tree: UInt[Array, '* 2**(d-1)']
-    split_tree: UInt[Array, '* 2**(d-1)']
+    leaf_tree: (
+        Float32[Array, '*batch_shape 2**d'] | Float32[Array, '*batch_shape k 2**d']
+    )
+    """The values in the leaves of the trees. This array can be dirty, i.e.,
+    unused nodes can have whatever value. It may have an additional axis
+    for multivariate leaves."""
+
+    var_tree: UInt[Array, '*batch_shape 2**(d-1)']
+    """The axes along which the decision nodes operate. This array can be
+    dirty but for the always unused node at index 0 which must be set to 0."""
+
+    split_tree: UInt[Array, '*batch_shape 2**(d-1)']
+    """The decision boundaries of the trees. The boundaries are open on the
+    right, i.e., a point belongs to the left child iff x < split. Whether a
+    node is a leaf is indicated by the corresponding 'split' element being
+    0. Unused nodes also have split set to 0. This array can't be dirty."""
 
 
-def make_tree(depth: int, dtype: DTypeLike) -> Shaped[Array, ' 2**{depth}']:
+def make_tree(
+    depth: int, dtype: DTypeLike, batch_shape: tuple[int, ...] = ()
+) -> Shaped[Array, '*batch_shape 2**{depth}']:
     """
     Make an array to represent a binary tree.
 
@@ -80,15 +85,19 @@ def make_tree(depth: int, dtype: DTypeLike) -> Shaped[Array, ' 2**{depth}']:
         node.
     dtype
         The dtype of the array.
+    batch_shape
+        The leading shape of the array, to represent multiple trees and/or
+        multivariate trees.
 
     Returns
     -------
     An array of zeroes with the appropriate shape.
     """
-    return jnp.zeros(2**depth, dtype)
+    shape = (*batch_shape, 2**depth)
+    return jnp.zeros(shape, dtype)
 
 
-def tree_depth(tree: Shaped[Array, '* 2**d']) -> int:
+def tree_depth(tree: Shaped[Array, '*batch_shape 2**d']) -> int:
     """
     Return the maximum depth of a tree.
 
@@ -106,10 +115,10 @@ def tree_depth(tree: Shaped[Array, '* 2**d']) -> int:
 
 
 def traverse_tree(
-    x: Real[Array, ' p'],
+    x: UInt[Array, ' p'],
     var_tree: UInt[Array, ' 2**(d-1)'],
     split_tree: UInt[Array, ' 2**(d-1)'],
-) -> Int32[Array, '']:
+) -> UInt[Array, '']:
     """
     Find the leaf where a point falls into.
 
@@ -148,15 +157,16 @@ def traverse_tree(
     return index
 
 
-@partial(vmap_nodoc, in_axes=(None, 0, 0))
+@jit
+@partial(jnp.vectorize, excluded=(0,), signature='(hts),(hts)->(n)')
 @partial(vmap_nodoc, in_axes=(1, None, None))
 def traverse_forest(
-    X: Real[Array, 'p n'],
-    var_trees: UInt[Array, 'm 2**(d-1)'],
-    split_trees: UInt[Array, 'm 2**(d-1)'],
-) -> Int32[Array, 'm n']:
+    X: UInt[Array, 'p n'],
+    var_trees: UInt[Array, '*forest_shape 2**(d-1)'],
+    split_trees: UInt[Array, '*forest_shape 2**(d-1)'],
+) -> UInt[Array, '*forest_shape n']:
     """
-    Find the leaves where points fall into.
+    Find the leaves where points falls into for each tree in a set.
 
     Parameters
     ----------
@@ -174,35 +184,59 @@ def traverse_forest(
     return traverse_tree(X, var_trees, split_trees)
 
 
+@partial(jit, static_argnames=('sum_batch_axis',))
 def evaluate_forest(
-    X: UInt[Array, 'p n'], trees: TreeHeaps, *, sum_trees: bool = True
-) -> Float32[Array, ' n'] | Float32[Array, 'm n']:
+    X: UInt[Array, 'p n'],
+    trees: TreeHeaps,
+    *,
+    sum_batch_axis: int | tuple[int, ...] = (),
+) -> (
+    Float32[Array, '*reduced_batch_size n'] | Float32[Array, '*reduced_batch_size k n']
+):
     """
-    Evaluate a ensemble of trees at an array of points.
+    Evaluate an ensemble of trees at an array of points.
 
     Parameters
     ----------
     X
         The coordinates to evaluate the trees at.
     trees
-        The tree heaps, with batch shape (m,).
-    sum_trees
-        Whether to sum the values across trees.
+        The trees.
+    sum_batch_axis
+        The batch axes to sum over. By default, no summation is performed.
+        Note that negative indices count from the end of the batch dimensions,
+        the core dimensions n and k can't be summed over by this function.
 
     Returns
     -------
     The (sum of) the values of the trees at the points in `X`.
     """
+    indices: UInt[Array, '*forest_shape n']
     indices = traverse_forest(X, trees.var_tree, trees.split_tree)
-    ntree, _ = trees.leaf_tree.shape
-    tree_index = jnp.arange(ntree, dtype=minimal_unsigned_dtype(ntree - 1))
-    leaves = trees.leaf_tree[tree_index[:, None], indices]
-    if sum_trees:
-        return jnp.sum(leaves, axis=0, dtype=jnp.float32)
-    # this sum suggests to swap the vmaps, but I think it's better for X
-    # copying to keep it that way
-    else:
-        return leaves
+
+    is_mv = trees.leaf_tree.ndim != trees.var_tree.ndim
+
+    bc_indices: UInt[Array, '*forest_shape n 1'] | UInt[Array, '*forest_shape 1 n 1']
+    bc_indices = indices[..., None, :, None] if is_mv else indices[..., None]
+
+    bc_leaf_tree: (
+        Float32[Array, '*forest_shape 1 tree_size']
+        | Float32[Array, '*forest_shape k 1 tree_size']
+    )
+    bc_leaf_tree = (
+        trees.leaf_tree[..., :, None, :] if is_mv else trees.leaf_tree[..., None, :]
+    )
+
+    bc_leaves: (
+        Float32[Array, '*forest_shape n 1'] | Float32[Array, '*forest_shape k n 1']
+    )
+    bc_leaves = jnp.take_along_axis(bc_leaf_tree, bc_indices, -1)
+
+    leaves: Float32[Array, '*forest_shape n'] | Float32[Array, '*forest_shape k n']
+    leaves = jnp.squeeze(bc_leaves, -1)
+
+    axis = normalize_axis_tuple(sum_batch_axis, trees.var_tree.ndim - 1)
+    return jnp.sum(leaves, axis=axis)
 
 
 def is_actual_leaf(
@@ -259,13 +293,13 @@ def is_leaves_parent(split_tree: UInt[Array, ' 2**(d-1)']) -> Bool[Array, ' 2**(
     # the 0-th item has split == 0, so it's not counted
 
 
-def tree_depths(tree_length: int) -> Int32[Array, ' {tree_length}']:
+def tree_depths(tree_size: int) -> Int32[Array, ' {tree_size}']:
     """
     Return the depth of each node in a binary tree.
 
     Parameters
     ----------
-    tree_length
+    tree_size
         The length of the tree array, i.e., 2 ** d.
 
     Returns
@@ -280,7 +314,7 @@ def tree_depths(tree_length: int) -> Int32[Array, ' {tree_length}']:
     """
     depths = []
     depth = 0
-    for i in range(tree_length):
+    for i in range(tree_size):
         if i == 2**depth:
             depth += 1
         depths.append(depth - 1)
@@ -288,7 +322,10 @@ def tree_depths(tree_length: int) -> Int32[Array, ' {tree_length}']:
     return jnp.array(depths, minimal_unsigned_dtype(max(depths)))
 
 
-def is_used(split_tree: UInt[Array, ' 2**(d-1)']) -> Bool[Array, ' 2**d']:
+@partial(jnp.vectorize, signature='(half_tree_size)->(tree_size)')
+def is_used(
+    split_tree: UInt[Array, '*batch_shape 2**(d-1)'],
+) -> Bool[Array, '*batch_shape 2**d']:
     """
     Return a mask indicating the used nodes in a tree.
 
@@ -308,7 +345,7 @@ def is_used(split_tree: UInt[Array, ' 2**(d-1)']) -> Bool[Array, ' 2**d']:
 
 
 @jit
-def forest_fill(split_tree: UInt[Array, 'num_trees 2**(d-1)']) -> Float32[Array, '']:
+def forest_fill(split_tree: UInt[Array, '*batch_shape 2**(d-1)']) -> Float32[Array, '']:
     """
     Return the fraction of used nodes in a set of trees.
 
@@ -321,36 +358,55 @@ def forest_fill(split_tree: UInt[Array, 'num_trees 2**(d-1)']) -> Float32[Array,
     -------
     Number of tree nodes over the maximum number that could be stored.
     """
-    num_trees, _ = split_tree.shape
-    used = jax.vmap(is_used)(split_tree)
+    used = is_used(split_tree)
     count = jnp.count_nonzero(used)
-    return count / (used.size - num_trees)
+    batch_size = split_tree.size // split_tree.shape[-1]
+    return count / (used.size - batch_size)
 
 
+@partial(jit, static_argnames=('p', 'sum_batch_axis'))
 def var_histogram(
-    p: int, var_tree: UInt[Array, '* 2**(d-1)'], split_tree: UInt[Array, '* 2**(d-1)']
-) -> Int32[Array, ' {p}']:
+    p: int,
+    var_tree: UInt[Array, '*batch_shape 2**(d-1)'],
+    split_tree: UInt[Array, '*batch_shape 2**(d-1)'],
+    *,
+    sum_batch_axis: int | tuple[int, ...] = (),
+) -> Int32[Array, '*reduced_batch_shape {p}']:
     """
     Count how many times each variable appears in a tree.
 
     Parameters
     ----------
     p
-        The number of variables (the maximum value that can occur in
-        `var_tree` is ``p - 1``).
+        The number of variables (the maximum value that can occur in `var_tree`
+        is ``p - 1``).
     var_tree
         The decision axes of the tree.
     split_tree
         The decision boundaries of the tree.
+    sum_batch_axis
+        The batch axes to sum over. By default, no summation is performed. Note
+        that negative indices count from the end of the batch dimensions, the
+        core dimension p can't be summed over by this function.
 
     Returns
     -------
-    The histogram of the variables used in the tree.
-
-    Notes
-    -----
-    If there are leading axes in the tree arrays (i.e., multiple trees), the
-    returned counts are cumulative over trees.
+    The histogram(s) of the variables used in the tree.
     """
     is_internal = split_tree.astype(bool)
-    return jnp.zeros(p, int).at[var_tree].add(is_internal)
+
+    def scatter_add(
+        var_tree: UInt[Array, '*summed_batch_axes half_tree_size'],
+        is_internal: Bool[Array, '*summed_batch_axes half_tree_size'],
+    ) -> Int32[Array, ' p']:
+        return jnp.zeros(p, int).at[var_tree].add(is_internal)
+
+    # vmap scatter_add over non-batched dims
+    batch_ndim = var_tree.ndim - 1
+    axes = normalize_axis_tuple(sum_batch_axis, batch_ndim)
+    for i in reversed(range(batch_ndim)):
+        neg_i = i - var_tree.ndim
+        if i not in axes:
+            scatter_add = vmap(scatter_add, in_axes=neg_i)
+
+    return scatter_add(var_tree, is_internal)

bartz/jaxext/__init__.py

@@ -1,6 +1,6 @@
 # bartz/src/bartz/jaxext/__init__.py
 #
-# Copyright (c) 2024-2025, Giacomo Petrillo
+# Copyright (c) 2024-2026, The Bartz Contributors
 #
 # This file is part of bartz.
 #
@@ -24,13 +24,23 @@
 
 """Additions to jax."""
 
-import functools
 import math
 from collections.abc import Sequence
+from contextlib import nullcontext
+from functools import partial
 
 import jax
+from jax import (
+    Device,
+    debug_key_reuse,
+    device_count,
+    ensure_compile_time_eval,
+    jit,
+    random,
+    vmap,
+)
 from jax import numpy as jnp
-from jax import random
+from jax.dtypes import prng_key
 from jax.lax import scan
 from jax.scipy.special import ndtr
 from jaxtyping import Array, Bool, Float32, Key, Scalar, Shaped
@@ -63,7 +73,7 @@ def minimal_unsigned_dtype(value):
     return jnp.uint64
 
 
-@functools.partial(jax.jit, static_argnums=(1,))
+@partial(jax.jit, static_argnums=(1,))
 def unique(
     x: Shaped[Array, ' _'], size: int, fill_value: Scalar
 ) -> tuple[Shaped[Array, ' {size}'], int]:
@@ -114,24 +124,42 @@ class split:
         The key to split.
     num
         The number of keys to split into.
+
+    Notes
+    -----
+    Unlike `jax.random.split`, this class supports a vector of keys as input. In
+    this case, it behaves as if everything had been vmapped over, so `keys.pop`
+    has an additional initial output dimension equal to the number of input
+    keys, and the deterministic dependency respects this axis.
     """
 
-    def __init__(self, key: Key[Array, ''], num: int = 2):
-        self._keys = random.split(key, num)
+    _keys: tuple[Key[Array, '*batch'], ...]
+    _num_used: int
+
+    def __init__(self, key: Key[Array, '*batch'], num: int = 2):
+        if key.ndim:
+            context = debug_key_reuse(False)
+        else:
+            context = nullcontext()
+        with context:
+            # jitted-vmapped key split seems to be triggering a false positive
+            # with key reuse checks
+            self._keys = _split_unpack(key, num)
+        self._num_used = 0
 
     def __len__(self):
-        return self._keys.size
+        return len(self._keys) - self._num_used
 
-    def pop(self, shape: int | tuple[int, ...] | None = None) -> Key[Array, '*']:
+    def pop(self, shape: int | tuple[int, ...] = ()) -> Key[Array, '*batch {shape}']:
         """
         Pop one or more keys from the list.
 
         Parameters
         ----------
         shape
-            The shape of the keys to pop. If `None`, a single key is popped.
-            If an integer, that many keys are popped. If a tuple, the keys are
-            reshaped to that shape.
+            The shape of the keys to pop. If empty (default), a single key is
+            popped and returned. If not empty, the popped key is split and
+            reshaped to the target shape.
 
         Returns
         -------
@@ -140,24 +168,41 @@ class split:
         Raises
         ------
         IndexError
-            If `shape` is larger than the number of keys left in the list.
-
-        Notes
-        -----
-        The keys are popped from the beginning of the list, so for example
-        ``list(keys.pop(2))`` is equivalent to ``[keys.pop(), keys.pop()]``.
+            If the list is empty.
         """
-        if shape is None:
-            shape = ()
-        elif not isinstance(shape, tuple):
-            shape = (shape,)
-        size_to_pop = math.prod(shape)
-        if size_to_pop > self._keys.size:
-            msg = f'Cannot pop {size_to_pop} keys from {self._keys.size} keys'
+        if len(self) == 0:
+            msg = 'No keys left to pop'
             raise IndexError(msg)
-        popped_keys = self._keys[:size_to_pop]
-        self._keys = self._keys[size_to_pop:]
-        return popped_keys.reshape(shape)
+        if not isinstance(shape, tuple):
+            shape = (shape,)
+        key = self._keys[self._num_used]
+        self._num_used += 1
+        if shape:
+            key = _split_shaped(key, shape)
+        return key
+
+
+@partial(jit, static_argnums=(1,))
+def _split_unpack(
+    key: Key[Array, '*batch'], num: int
+) -> tuple[Key[Array, '*batch'], ...]:
+    if key.ndim == 0:
+        keys = random.split(key, num)
+    elif key.ndim == 1:
+        keys = vmap(random.split, in_axes=(0, None), out_axes=1)(key, num)
+    return tuple(keys)
+
+
+@partial(jit, static_argnums=(1,))
+def _split_shaped(
+    key: Key[Array, '*batch'], shape: tuple[int, ...]
+) -> Key[Array, '*batch {shape}']:
+    num = math.prod(shape)
+    if key.ndim == 0:
+        keys = random.split(key, num)
+    elif key.ndim == 1:
+        keys = vmap(random.split, in_axes=(0, None))(key, num)
+    return keys.reshape(*key.shape, *shape)
 
 
 def truncated_normal_onesided(
@@ -165,6 +210,8 @@ def truncated_normal_onesided(
     shape: Sequence[int],
     upper: Bool[Array, '*'],
     bound: Float32[Array, '*'],
+    *,
+    clip: bool = True,
 ) -> Float32[Array, '*']:
     """
     Sample from a one-sided truncated standard normal distribution.
@@ -179,6 +226,9 @@ def truncated_normal_onesided(
         True for (-∞, bound], False for [bound, ∞).
     bound
         The truncation boundary.
+    clip
+        Whether to clip the truncated uniform samples to (0, 1) before
+        transforming them to truncated normal. Intended for debugging purposes.
 
     Returns
     -------
@@ -209,5 +259,29 @@ def truncated_normal_onesided(
     left_u = scale * (1 - u)  # ~ uniform in (0, ndtr(±bound)]
     right_u = shift + scale * u  # ~ uniform in [ndtr(∓bound), 1)
     truncated_u = jnp.where(upper ^ bound_pos, left_u, right_u)
+    if clip:
+        # on gpu the accuracy is lower and sometimes u can reach the boundaries
+        zero = jnp.zeros((), truncated_u.dtype)
+        one = jnp.ones((), truncated_u.dtype)
+        truncated_u = jnp.clip(
+            truncated_u, jnp.nextafter(zero, one), jnp.nextafter(one, zero)
+        )
     truncated_norm = ndtri(truncated_u)
     return jnp.where(bound_pos, -truncated_norm, truncated_norm)
+
+
+def get_default_device() -> Device:
+    """Get the current default JAX device."""
+    with ensure_compile_time_eval():
+        return jnp.zeros(()).device
+
+
+def get_device_count() -> int:
+    """Get the number of available devices on the default platform."""
+    device = get_default_device()
+    return device_count(device.platform)
+
+
+def is_key(x: object) -> bool:
+    """Determine if `x` is a jax random key."""
+    return isinstance(x, Array) and jnp.issubdtype(x.dtype, prng_key)