optax.tree_utils.tree_full_like

Contents

optax.tree_utils.tree_full_like#

optax.tree_utils.tree_full_like(
tree: Any,
fill_value: jax.typing.ArrayLike,
dtype: str | type[Any] | dtype | SupportsDType | None = None,
) → Any[source]#

Creates an identical tree where all tensors are filled with fill_value.

Parameters:

  • tree – pytree.

  • fill_value – the fill value for all tensors in the tree.

  • dtype – optional dtype to use for the tensors in the tree.

Returns:

an tree with the same structure as tree.

Contents