optax.tree_utils.tree_ones_like

optax.tree_utils.tree_ones_like(

tree: Any,

dtype: str | type[Any] | dtype | SupportsDType | None = None,

) → Any

Creates an all-ones tree with the same structure.

Parameters:

• tree – pytree.

• dtype – optional dtype to use for the tree of ones.

Returns:

an all-ones tree with the same structure as tree.