tree (Quantity]) – Input PyTree.
idx (int | slice | Array | ndarray | bool | number | bool | float | complex | Quantity) – 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.