braintools.tree.idx

Contents

braintools.tree.idx#

braintools.tree.idx(tree, idx, is_leaf=<function is_quantity>)[source]#

Index every leaf of a PyTree.

Parameters:

  • tree (Quantity]) – Input PyTree.

  • idx (int, slice or array_like) – Indices used to index each leaf x as x[idx].

  • is_leaf (Callable[[Any], bool] | None) – Predicate to treat a node as a leaf during traversal. Defaults to u.math.is_quantity.

Returns:

PyTree with leaves indexed by idx.

Return type:

PyTree

Contents