braintools.tree.squared_norm

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

Sum of squares of all elements in a PyTree.

Parameters:

• tree (Quantity]) – Input PyTree.

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

Returns:

Scalar representing sum_i sum(x_i**2) across all leaves x_i.

Return type:

Array

See also

dot

Inner product across leaves.

Notes

Equivalent to dot(tree, tree).