braintools.tree.scale

Contents

braintools.tree.scale#

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

Multiply every leaf array in a PyTree by a value.

Parameters:
  • tree (Quantity]) – Input PyTree whose leaves are multiplied.

  • x (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Value broadcast and multiplied with each leaf of tree.

  • is_leaf (Callable[[Any], bool] | None) – Predicate to treat a node as a leaf during traversal. Defaults to u.math.is_quantity so unit-quantities are considered leaves.

Returns:

PyTree with each leaf multiplied by x (with broadcasting).

Return type:

PyTree

See also

mul

Elementwise multiplication with a PyTree or value.

add

Elementwise addition with a PyTree or value.

Notes

  • Broadcasting follows the semantics of the underlying array type (e.g., JAX/NumPy). The structure of tree is preserved.