Coverage for src/bartz/grove/_grove.py: 99%
199 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 11:03 +0000
1# bartz/src/bartz/grove/_grove.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"""Functions to create and manipulate binary decision trees."""
27import math
28from dataclasses import fields
29from functools import partial
30from typing import Literal, Protocol, runtime_checkable
32from equinox import tree_at
33from jax import numpy as jnp
34from jax import vmap
35from jaxtyping import Array, Bool, Float, Float32, Int32, Shaped, UInt
36from numpy.lib.array_utils import normalize_axis_tuple
38from bartz._jaxext import (
39 Module,
40 autobatch,
41 field,
42 jit,
43 minimal_unsigned_dtype,
44 vmap_nodoc,
45)
48@runtime_checkable
49class TreeHeaps(Protocol):
50 """A protocol for dataclasses that represent trees.
52 A tree is represented with arrays as a heap. The root node is at index 1.
53 The children nodes of a node at index :math:`i` are at indices :math:`2i`
54 (left child) and :math:`2i + 1` (right child). The array element at index 0
55 is unused.
57 Since the nodes at the bottom can only be leaves and not decision nodes,
58 `var_tree` and `split_tree` are half as long as `leaf_tree`.
60 Arrays may have additional initial axes to represent multiple trees.
61 """
63 leaf_tree: (
64 Float[Array, '*batch_shape 2*half_tree_size']
65 | Float[Array, '*batch_shape k 2*half_tree_size']
66 )
67 """The values in the leaves of the trees. This array can be dirty, i.e.,
68 unused nodes can have whatever value. It may have an additional axis
69 for multivariate leaves. The dtype may be narrower than float32."""
71 var_tree: UInt[Array, '*batch_shape half_tree_size']
72 """The axes along which the decision nodes operate. This array can be
73 dirty but for the always unused node at index 0 which must be set to 0."""
75 split_tree: UInt[Array, '*batch_shape half_tree_size']
76 """The decision boundaries of the trees. The boundaries are open on the
77 right, i.e., a point belongs to the left child iff x < split. Whether a
78 node is a leaf is indicated by the corresponding 'split' element being
79 0. Unused nodes also have split set to 0. This array can't be dirty."""
81 leaf_unit: Float32[Array, ''] | Float32[Array, ' k']
82 """The storage unit of the leaf values. The leaf values in data units are
83 ``leaf_unit * leaf_tree``."""
85 offset: Float32[Array, ''] | Float32[Array, ' k']
86 """Constant shift added to the function represented by the trees, which is
87 ``offset + leaf_unit * (sum of leaf values over trees)``."""
90def is_multivariate(trees: TreeHeaps) -> bool:
91 """
92 Return whether the trees have vector-valued leaves.
94 Parameters
95 ----------
96 trees
97 The trees to inspect.
99 Returns
100 -------
101 Whether the leaves are vector-valued (an extra `k` axis on `leaf_tree`).
102 """
103 return trees.leaf_tree.ndim > trees.var_tree.ndim
106class TreesTrace(Module):
107 """Implementation of `bartz.grove.TreeHeaps` for an MCMC trace."""
109 # `var_tree`/`split_tree` are declared before `leaf_tree` so their single
110 # (union-free) annotations bind the variadic `*batch_shape` axis first;
111 # otherwise the runtime typechecker (which evaluates union members in a
112 # hash-randomized order) can mis-bind it against the `k` axis of
113 # `leaf_tree`'s union for a multivariate tree (the layouts are
114 # rank-ambiguous). See `bartz.mcmcstep._state.Forest`. The leaf-bearing axis
115 # is `2*half_tree_size` rather than `tree_size`, so the half-of-leaf
116 # relationship is still checked here: `half_tree_size` is bound first by the
117 # anchors, then `leaf_tree` is checked against twice it.
118 var_tree: UInt[Array, '*batch_shape half_tree_size']
119 """The axes along which the decision nodes operate. This array can be
120 dirty but for the always unused node at index 0 which must be set to 0."""
122 split_tree: UInt[Array, '*batch_shape half_tree_size']
123 """The decision boundaries of the trees. The boundaries are open on the
124 right, i.e., a point belongs to the left child iff x < split. Whether a
125 node is a leaf is indicated by the corresponding 'split' element being
126 0. Unused nodes also have split set to 0. This array can't be dirty."""
128 leaf_tree: (
129 Float[Array, '*batch_shape 2*half_tree_size']
130 | Float[Array, '*batch_shape k 2*half_tree_size']
131 )
132 """The values in the leaves of the trees. This array can be dirty, i.e.,
133 unused nodes can have whatever value. It may have an additional axis
134 for multivariate leaves."""
136 leaf_unit: Float32[Array, ''] | Float32[Array, ' k'] = field(
137 default_factory=lambda: jnp.float32(1.0)
138 )
139 offset: Float32[Array, ''] | Float32[Array, ' k'] = field(
140 default_factory=lambda: jnp.float32(0.0)
141 )
143 def axes_from_dataclass(self, obj: TreeHeaps) -> 'TreesTrace':
144 """Project the per-field vmap axis specs of `obj` onto this template.
146 `self` supplies the (array) pytree; the same-named fields of `obj`
147 (axis specs, i.e. ints or `None`) replace its leaves. Built with
148 `equinox.tree_at`, which bypasses the type-checked `__init__`, so the
149 deliberately off-type axis values are allowed.
150 """
151 names = [f.name for f in fields(type(self))]
152 return tree_at(
153 lambda t: [getattr(t, name) for name in names],
154 self,
155 [getattr(obj, name) for name in names],
156 )
159def tree_depth(tree: Shaped[Array, '*batch_shape tree_size']) -> int:
160 """
161 Return the maximum depth of a tree.
163 Parameters
164 ----------
165 tree
166 A tree array like those in a `TreeHeaps`. If the array is ND, the tree
167 structure is assumed to be along the last axis.
169 Returns
170 -------
171 The maximum depth of the tree.
172 """
173 return round(math.log2(tree.shape[-1]))
176def traverse_tree(
177 x: UInt[Array, ' p'],
178 var_tree: UInt[Array, ' half_tree_size'],
179 split_tree: UInt[Array, ' half_tree_size'],
180) -> UInt[Array, '']:
181 """
182 Find the leaf where a point falls into.
184 Parameters
185 ----------
186 x
187 The coordinates to evaluate the tree at.
188 var_tree
189 The decision axes of the tree.
190 split_tree
191 The decision boundaries of the tree.
193 Returns
194 -------
195 The index of the leaf.
196 """
197 leaf_found = jnp.zeros((), bool)
198 index = jnp.ones((), minimal_unsigned_dtype(2 * var_tree.size - 1))
200 # the depth is a small static integer, so a plain python loop is equivalent
201 # to (and clearer than) a fully-unrolled lax.scan
202 for _ in range(tree_depth(var_tree)):
203 split = split_tree[index]
204 var = var_tree[index]
206 leaf_found |= split == 0
207 child_index = (index << 1) + (x[var] >= split)
208 index = jnp.where(leaf_found, index, child_index)
210 return index
213@jit
214def traverse_forest(
215 X: UInt[Array, 'p n'],
216 var_trees: UInt[Array, '*forest_shape half_tree_size'],
217 split_trees: UInt[Array, '*forest_shape half_tree_size'],
218) -> UInt[Array, '*forest_shape n']:
219 """
220 Find the leaves where points falls into for each tree in a set.
222 Parameters
223 ----------
224 X
225 The coordinates to evaluate the trees at.
226 var_trees
227 The decision axes of the trees.
228 split_trees
229 The decision boundaries of the trees.
231 Returns
232 -------
233 The indices of the leaves.
234 """
235 return _traverse_forest(X, var_trees, split_trees)
238@partial(jnp.vectorize, excluded=(0,), signature='(hts),(hts)->(n)')
239@partial(vmap_nodoc, in_axes=(1, None, None))
240def _traverse_forest(
241 X: UInt[Array, ' p'],
242 var_trees: UInt[Array, ' half_tree_size'],
243 split_trees: UInt[Array, ' half_tree_size'],
244) -> UInt[Array, '']:
245 """Implement `traverse_forest`."""
246 return traverse_tree(X, var_trees, split_trees)
249@jit(static_argnames=('sum_batch_axis',))
250def evaluate_forest(
251 X: UInt[Array, 'p n'],
252 trees: TreeHeaps,
253 *,
254 sum_batch_axis: int | tuple[int, ...] = (),
255) -> (
256 Float32[Array, '*reduced_batch_size n'] | Float32[Array, '*reduced_batch_size k n']
257):
258 """
259 Evaluate an ensemble of trees at an array of points.
261 Parameters
262 ----------
263 X
264 The coordinates to evaluate the trees at.
265 trees
266 The trees.
267 sum_batch_axis
268 The batch axes to sum over. By default, no summation is performed.
269 Note that negative indices count from the end of the batch dimensions,
270 the core dimensions n and k can't be summed over by this function.
272 Returns
273 -------
274 The (sum of) the values of the trees at the points in `X`, in float32.
276 Notes
277 -----
278 The leaf values are multiplied by ``trees.leaf_unit``, so the result is in
279 data units. ``trees.offset`` is not added: it does not distribute over the
280 per-tree sum, so the caller adds it once to the total.
281 """
282 indices: UInt[Array, '*forest_shape n']
283 indices = traverse_forest(X, trees.var_tree, trees.split_tree)
285 is_mv = is_multivariate(trees)
287 bc_indices: UInt[Array, '*forest_shape n 1'] | UInt[Array, '*forest_shape 1 n 1']
288 bc_indices = indices[..., None, :, None] if is_mv else indices[..., None]
290 bc_leaf_tree: (
291 Float[Array, '*forest_shape 1 tree_size']
292 | Float[Array, '*forest_shape k 1 tree_size']
293 )
294 bc_leaf_tree = (
295 trees.leaf_tree[..., :, None, :] if is_mv else trees.leaf_tree[..., None, :]
296 )
298 bc_leaves: Float[Array, '*forest_shape n 1'] | Float[Array, '*forest_shape k n 1']
299 bc_leaves = jnp.take_along_axis(bc_leaf_tree, bc_indices, -1)
301 leaves: Float[Array, '*forest_shape n'] | Float[Array, '*forest_shape k n']
302 leaves = jnp.squeeze(bc_leaves, -1)
304 # sum in the storage dtype with a float32 accumulator, then convert to
305 # data units with the float32 scale
306 axis = normalize_axis_tuple(sum_batch_axis, trees.var_tree.ndim - 1)
307 return trees.leaf_unit[..., None] * jnp.sum(leaves, axis=axis, dtype=jnp.float32)
310def is_actual_leaf(
311 split_tree: UInt[Array, ' half_tree_size'], *, add_bottom_level: bool = False
312) -> Bool[Array, ' half_tree_size'] | Bool[Array, ' 2*half_tree_size']:
313 """
314 Return a mask indicating the leaf nodes in a tree.
316 Parameters
317 ----------
318 split_tree
319 The splitting points of the tree.
320 add_bottom_level
321 If True, the bottom level of the tree is also considered.
323 Returns
324 -------
325 The mask marking the leaf nodes. Length doubled if `add_bottom_level` is True.
326 """
327 size = split_tree.size
328 is_leaf = split_tree == 0
329 if add_bottom_level:
330 size *= 2
331 is_leaf = jnp.concatenate([is_leaf, jnp.ones_like(is_leaf)])
332 index = jnp.arange(size, dtype=minimal_unsigned_dtype(size - 1))
333 parent_index = index >> 1
334 parent_nonleaf = split_tree[parent_index].astype(bool)
335 parent_nonleaf = parent_nonleaf.at[1].set(True)
336 return is_leaf & parent_nonleaf
339def is_leaves_parent(
340 split_tree: UInt[Array, ' half_tree_size'],
341) -> Bool[Array, ' half_tree_size']:
342 """
343 Return a mask indicating the nodes with leaf (and only leaf) children.
345 Parameters
346 ----------
347 split_tree
348 The decision boundaries of the tree.
350 Returns
351 -------
352 The mask indicating which nodes have leaf children.
353 """
354 index = jnp.arange(
355 split_tree.size, dtype=minimal_unsigned_dtype(2 * split_tree.size - 1)
356 )
357 left_index = index << 1 # left child
358 right_index = left_index + 1 # right child
359 left_leaf = split_tree.at[left_index].get(mode='fill', fill_value=0) == 0
360 right_leaf = split_tree.at[right_index].get(mode='fill', fill_value=0) == 0
361 is_not_leaf = split_tree.astype(bool)
362 return is_not_leaf & left_leaf & right_leaf
363 # the 0-th item has split == 0, so it's not counted
366def tree_depths(tree_size: int) -> UInt[Array, ' {tree_size}']:
367 """
368 Return the depth of each node in a binary tree.
370 Parameters
371 ----------
372 tree_size
373 The length of the tree array, i.e., 2 ** d.
375 Returns
376 -------
377 The depth of each node.
379 Notes
380 -----
381 The root node (index 1) has depth 0. The depth is the position of the most
382 significant non-zero bit in the index. The first element (the unused node)
383 is marked as depth 0.
384 """
385 depths = []
386 depth = 0
387 for i in range(tree_size):
388 if i == 2**depth:
389 depth += 1
390 depths.append(depth - 1)
391 depths[0] = 0
392 return jnp.array(depths, minimal_unsigned_dtype(max(depths)))
395@jit
396def forest_mean_leaves(
397 split_tree: UInt[Array, '*batch_shape half_tree_size'],
398) -> Float32[Array, '']:
399 """
400 Return the average number of leaves per tree in a set of trees.
402 Parameters
403 ----------
404 split_tree
405 The decision boundaries of the trees.
407 Returns
408 -------
409 The mean number of leaves across the trees.
410 """
411 # a tree with k internal nodes (the nonzero entries of split_tree) has k + 1
412 # leaves; the maximum possible is split_tree.shape[-1]
413 num_internal = jnp.count_nonzero(split_tree, axis=-1)
414 return (num_internal + 1).mean()
417@jit(static_argnames=('p', 'sum_batch_axis'))
418def var_histogram(
419 p: int,
420 var_tree: UInt[Array, '*batch_shape half_tree_size'],
421 split_tree: UInt[Array, '*batch_shape half_tree_size'],
422 *,
423 sum_batch_axis: int | tuple[int, ...] = (),
424) -> Int32[Array, '*reduced_batch_shape {p}']:
425 """
426 Count how many times each variable appears in a tree.
428 Parameters
429 ----------
430 p
431 The number of variables (the maximum value that can occur in `var_tree`
432 is ``p - 1``).
433 var_tree
434 The decision axes of the tree.
435 split_tree
436 The decision boundaries of the tree.
437 sum_batch_axis
438 The batch axes to sum over. By default, no summation is performed. Note
439 that negative indices count from the end of the batch dimensions, the
440 core dimension p can't be summed over by this function.
442 Returns
443 -------
444 The histogram(s) of the variables used in the tree.
445 """
446 is_internal = split_tree.astype(bool)
448 def scatter_add(
449 var_tree: UInt[Array, '*summed_batch_axes half_tree_size'],
450 is_internal: Bool[Array, '*summed_batch_axes half_tree_size'],
451 ) -> Int32[Array, ' p']:
452 return jnp.zeros(p, int).at[var_tree].add(is_internal)
454 # vmap scatter_add over non-batched dims
455 batch_ndim = var_tree.ndim - 1
456 axes = normalize_axis_tuple(sum_batch_axis, batch_ndim)
457 for i in reversed(range(batch_ndim)):
458 neg_i = i - var_tree.ndim
459 if i not in axes:
460 scatter_add = vmap(scatter_add, in_axes=neg_i)
462 return scatter_add(var_tree, is_internal)
465def _format_values(values: Float[Array, ''] | Float[Array, ' k']) -> str:
466 """Format a scalar or vector to the decimal precision of its dtype."""
467 ndigits = jnp.finfo(values.dtype).precision
468 if values.ndim:
469 return '[' + ', '.join(f'{v:#.{ndigits}g}' for v in values) + ']'
470 else:
471 return f'{values:#.{ndigits}g}'
474def _format_leaf(tree: TreeHeaps, index: int) -> str:
475 """Format the leaf value at `index` as ``<scale> * <leaf>``."""
476 unit = _format_values(tree.leaf_unit)
477 leaf = _format_values(tree.leaf_tree[..., index])
478 return f'{unit} * {leaf}'
481def format_tree(tree: TreeHeaps, *, print_all: bool = False) -> str:
482 """Convert a tree to a human-readable string.
484 Parameters
485 ----------
486 tree
487 A single tree to format.
488 print_all
489 If `True`, also print the contents of unused node slots in the arrays.
491 Returns
492 -------
493 A string representation of the tree.
494 """
495 tee = '├──'
496 corner = '└──'
497 join = '│ '
498 space = ' '
499 down = '┐'
500 bottom = '╢' # '┨' #
502 *_, tree_size = tree.leaf_tree.shape
504 def traverse_tree(
505 lines: list[str],
506 index: int,
507 depth: int,
508 indent: str,
509 first_indent: str,
510 next_indent: str,
511 unused: bool,
512 ) -> None:
513 if index >= tree_size:
514 return
516 var: int = tree.var_tree.at[index].get(mode='fill', fill_value=0).item()
517 split: int = tree.split_tree.at[index].get(mode='fill', fill_value=0).item()
519 is_leaf = split == 0
520 left_child = 2 * index
521 right_child = 2 * index + 1
523 if print_all:
524 if unused:
525 category = 'unused'
526 elif is_leaf:
527 category = 'leaf'
528 else:
529 category = 'decision'
530 node_str = f'{category}({var}, {split}, {_format_leaf(tree, index)})'
531 else:
532 assert not unused
533 if is_leaf:
534 node_str = _format_leaf(tree, index)
535 else:
536 node_str = f'x{var} < {split}'
538 if not is_leaf or (print_all and left_child < tree_size):
539 link = down
540 elif not print_all and left_child >= tree_size:
541 link = bottom
542 else:
543 link = ' '
545 max_number = tree_size - 1
546 ndigits = len(str(max_number))
547 number = str(index).rjust(ndigits)
549 lines.append(f' {number} {indent}{first_indent}{link}{node_str}')
551 indent += next_indent
552 unused = unused or is_leaf
554 if unused and not print_all:
555 return
557 traverse_tree(lines, left_child, depth + 1, indent, tee, join, unused)
558 traverse_tree(lines, right_child, depth + 1, indent, corner, space, unused)
560 lines = []
561 traverse_tree(lines, 1, 0, '', '', '', False)
562 return '\n'.join(lines)
565def tree_actual_depth(split_tree: UInt[Array, ' half_tree_size']) -> UInt[Array, '']:
566 """Measure the depth of the tree.
568 Parameters
569 ----------
570 split_tree
571 The cutpoints of the decision rules.
573 Returns
574 -------
575 The depth of the deepest leaf in the tree. The root is at depth 0.
576 """
577 # this could be done just with split_tree != 0
578 is_leaf = is_actual_leaf(split_tree, add_bottom_level=True)
579 depth = tree_depths(is_leaf.size)
580 depth = jnp.where(is_leaf, depth, 0)
581 return jnp.max(depth)
584@jit
585@partial(jnp.vectorize, signature='(nt,hts)->(d)')
586def forest_depth_distr(
587 split_tree: UInt[Array, '*batch_shape num_trees half_tree_size'],
588) -> Int32[Array, '*batch_shape d']:
589 """Histogram the depths of a set of trees.
591 Parameters
592 ----------
593 split_tree
594 The cutpoints of the decision rules of the trees.
596 Returns
597 -------
598 An integer vector where the i-th element counts how many trees have depth i.
599 """
600 depth = tree_depth(split_tree) + 1
601 depths = vmap(tree_actual_depth)(split_tree)
602 return jnp.bincount(depths, length=depth)
605@jit(static_argnames=('node_type', 'sum_batch_axis'))
606def points_per_node_distr(
607 X: UInt[Array, 'p n'],
608 var_tree: UInt[Array, '*batch_shape half_tree_size'],
609 split_tree: UInt[Array, '*batch_shape half_tree_size'],
610 node_type: Literal['leaf', 'leaf-parent'],
611 *,
612 sum_batch_axis: int | tuple[int, ...] = (),
613) -> Int32[Array, '*reduced_batch_shape n+1']:
614 """Histogram points-per-node counts in a set of trees.
616 Count how many nodes in a tree select each possible amount of points,
617 over a certain subset of nodes.
619 Parameters
620 ----------
621 X
622 The set of points to count.
623 var_tree
624 The variables of the decision rules.
625 split_tree
626 The cutpoints of the decision rules.
627 node_type
628 The type of nodes to consider. Can be:
630 'leaf'
631 Count only leaf nodes.
632 'leaf-parent'
633 Count only parent-of-leaf nodes.
634 sum_batch_axis
635 Aggregate the histogram over these batch axes, counting how many nodes
636 have each possible amount of points over subsets of trees instead of
637 in each tree separately.
639 Returns
640 -------
641 A vector where the i-th element counts how many nodes have i points.
642 """
643 batch_ndim = var_tree.ndim - 1
644 axes = normalize_axis_tuple(sum_batch_axis, batch_ndim)
646 def func(
647 var_tree: UInt[Array, '*batch_shape half_tree_size'],
648 split_tree: UInt[Array, '*batch_shape half_tree_size'],
649 ) -> Int32[Array, '*reduced_batch_shape n_plus_1']:
650 indices: UInt[Array, '*batch_shape n']
651 indices = traverse_forest(X, var_tree, split_tree)
653 @partial(jnp.vectorize, signature='(hts),(n)->(ts_or_hts),(ts_or_hts)')
654 def count_points(
655 split_tree: UInt[Array, '*batch_shape half_tree_size'],
656 indices: UInt[Array, '*batch_shape n'],
657 ) -> (
658 tuple[
659 Int32[Array, '*batch_shape 2*half_tree_size'],
660 Bool[Array, '*batch_shape 2*half_tree_size'],
661 ]
662 | tuple[
663 Int32[Array, '*batch_shape half_tree_size'],
664 Bool[Array, '*batch_shape half_tree_size'],
665 ]
666 ):
667 if node_type == 'leaf-parent':
668 indices >>= 1
669 predicate = is_leaves_parent(split_tree)
670 elif node_type == 'leaf': 670 ↛ 673line 670 didn't jump to line 673 because the condition on line 670 was always true
671 predicate = is_actual_leaf(split_tree, add_bottom_level=True)
672 else:
673 raise ValueError(node_type)
674 count_tree = jnp.zeros(predicate.size, int).at[indices].add(1).at[0].set(0)
675 return count_tree, predicate
677 count_tree, predicate = count_points(split_tree, indices)
679 def count_nodes(
680 count_tree: Int32[Array, '*summed_batch_axes half_tree_size'],
681 predicate: Bool[Array, '*summed_batch_axes half_tree_size'],
682 ) -> Int32[Array, ' n_plus_1']:
683 return jnp.zeros(X.shape[1] + 1, int).at[count_tree].add(predicate)
685 # vmap count_nodes over non-batched dims
686 for i in reversed(range(batch_ndim)):
687 neg_i = i - var_tree.ndim
688 if i not in axes:
689 count_nodes = vmap(count_nodes, in_axes=neg_i)
691 return count_nodes(count_tree, predicate)
693 # automatically batch over all batch dimensions
694 max_io_nbytes = 2**27 # 128 MiB
695 out_dim_shift = len(axes)
696 batched_func = func
697 for i in reversed(range(batch_ndim)):
698 if i in axes:
699 out_dim_shift -= 1
700 else:
701 batched_func = autobatch(batched_func, max_io_nbytes, i, i - out_dim_shift)
702 assert out_dim_shift == 0
704 return batched_func(var_tree, split_tree)