Serialization

klax.text_serialize_filter_spec(f, x)

Filter specification for serializing a leaf to text.

PARAMETER	DESCRIPTION
f	File-like object to write to. TYPE: BinaryIO
x	The leaf to save in the file. TYPE: Any

Example

Serializing a model to a text file.

>>> import equinox as eqx
>>> import jax.numpy as jnp
>>> import klax
>>>
>>> tree = (jnp.array([1, 2, 3]), [3, 4, 5])
>>> eqx.tree_serialize_leaves(
...     "some_txt_file.txt",
...     tree,
...     filter_spec=klax.text_serialize_filter_spec
... )

klax.text_deserialize_filter_spec(f, x)

Filter specification for deserializing a leaf from text.

This function can be used to deserialized leafs that have been serialized using klax.text_serialize_filter_spec.

PARAMETER	DESCRIPTION
f	File-like object to read from. TYPE: BinaryIO
x	The leaf for which to load data from the file. TYPE: Any

Example

>>> import equinox as eqx
>>> import jax.numpy as jnp
>>> import klax
>>>
>>> tree = (jnp.array([1, 2, 3]), [3, 4, 5])
>>> eqx.tree_serialize_leaves(
...     "some_txt_file.txt",
...     tree,
...     filter_spec=klax.text_serialize_filter_spec
... )
>>> loaded_tree = eqx.tree_deserialize_leaves(
...     "some_txt_file.txt",
...     tree,
...     filter_spec=klax.text_deserialize_filter_spec
... )