optax.tree_utils.tree_ones_like#
- optax.tree_utils.tree_ones_like(tree: Any, dtype: str | type[Any] | dtype | SupportsDType | None = None) Any[source]#
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.