braintools.tree.sum

braintools.tree.sum(tree, is_leaf=<function is_quantity>)

Sum all elements across every leaf in a PyTree.

Parameters:

• tree (Quantity]) – Input PyTree whose elements are reduced.

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

Returns:

Scalar sum across all leaves and their elements.

Return type:

Array

See also

dot

Inner product across leaves.

squared_norm

Sum of squares across leaves.