optax.tree_utils.tree_add_scale

optax.tree_utils.tree_add_scale(

tree_x: Any,

scalar: jax.typing.ArrayLike,

tree_y: Any,

) → Any

Add two trees, where the second tree is scaled by a scalar.

In infix notation, the function performs out = tree_x + scalar * tree_y.

Parameters:

• tree_x – first pytree.

• scalar – scalar value.

• tree_y – second pytree.

Returns:

a pytree with the same structure as tree_x and tree_y.