optax.tree_utils.tree_where

optax.tree_utils.tree_where(condition, tree_x, tree_y)

Select tree_x values if condition is true else tree_y values.

Parameters:

• condition – boolean specifying which values to select from tree x or tree_y

• tree_x – pytree chosen if condition is True

• tree_y – pytree chosen if condition is False

Returns:

tree_x or tree_y depending on condition.