Ctrl+K

etree.spec_like

Contents

etree.spec_like#

etils.etree.spec_like(tree: etils.enp.array_types.typing.Array | list[etils.etree.typing._T | list[Tree] | tuple[Tree, ...] | dict[str, Tree]] | tuple[etils.etree.typing._T | list[Tree] | tuple[Tree, ...] | dict[str, Tree], ...] | dict[str, etils.etree.typing._T | list[Tree] | tuple[Tree, ...] | dict[str, Tree]], *, ignore_other: bool = True) → etils.enp.array_spec.ArraySpec | list[etils.etree.typing._T | list[Tree] | tuple[Tree, ...] | dict[str, Tree]] | tuple[etils.etree.typing._T | list[Tree] | tuple[Tree, ...] | dict[str, Tree], ...] | dict[str, etils.etree.typing._T | list[Tree] | tuple[Tree, ...] | dict[str, Tree]][source]#

Inspect a tree of array, works with any array type.

Example:

model = MyModel()
variables = model.init(jax.random.PRNGKey(0), x)

# Inspect the `variables` tree structures
print(etree.spec_like(variables))
Parameters:

  • tree – The tree of array

  • ignore_other – If True, non-array are forwarded as-is.

Returns:

The tree of enp.ArraySpec.

Contents