optax.tree_utils.tree_clip

Contents

optax.tree_utils.tree_clip#

optax.tree_utils.tree_clip(
tree: Any,
min_value: TypeAliasForwardRef('jax.typing.ArrayLike') | None = None,
max_value: TypeAliasForwardRef('jax.typing.ArrayLike') | None = None,
) → Any[source]#

Creates an identical tree where all tensors are clipped to [min, max].

Parameters:

  • tree – pytree.

  • min_value – optional minimal value to clip all tensors to. If None (default) then result will not be clipped to any minimum value.

  • max_value – optional maximal value to clip all tensors to. If None (default) then result will not be clipped to any maximum value.

Returns:

a tree with the same structure as tree.

Added in version 0.2.3.

Contents