optax.tree_utils.tree_add

optax.tree_utils.tree_add(

tree_x: Any,

tree_y: Any,

*other_trees: Any,

) → Any

Add two (or more) pytrees.

Parameters:

• tree_x – first pytree.

• tree_y – second pytree.

• *other_trees – optional other trees to add

Returns:

the sum of the two (or more) pytrees.

Changed in version 0.2.1: Added optional *other_trees argument.