🔗 Common API

class serket.Sequential(*a, **k)

A sequential container for layers.

Parameters:

layers – a tuple or a list of layers. if a list is passed, it will be casted to a tuple to maintain immutable behavior.

Example

>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> import serket as sk
>>> layers = sk.Sequential(lambda x: x + 1, lambda x: x * 2)
>>> print(layers(jnp.array([1, 2, 3]), key=jr.key(0)))
[4 6 8]

Note

Layer might be a function or a class with a __call__ method, additionally it might have a key argument for random number generation.

__call__(input, *, key=None)

Call self as a function.

Parameters:

• input (Array)

• key (Optional[Array])

Return type:

Array

serket.tree_state(tree, **kwargs)

Build state for a tree of layers.

Some layers require state to be initialized before training. For example, nn.BatchNorm layers requires running_mean and running_var to be initialized before training. This function initializes the state for a tree of layers, based on the layer defined state rule using tree_state.def_state.

tree_state() objective is to provide a simple and consistent way to initialize state for a tree of layers. Specifically, it provides a way to separate the state initialization logic from the layer definition. This allows for a more clear separation of concerns, and makes it easier to define new layers.

Parameters:

• tree (TypeVar(T)) – A tree of layers.

• kwargs – Keyword arguments to pass to the state initialization rule. of the tree layers.

Return type:

TypeVar(T)

Returns:

A tree of state leaves if it has state, otherwise NoState placeholder.

Note

To define a state initialization rule for a custom layer, use the decorator tree_state.def_state() on a function that accepts the layer as the first argument, for any additional arguments, use keyword only arguments.

>>> import jax
>>> import serket as sk
>>> class LayerWithState(sk.TreeClass):
...    pass
>>> # state function accept the `layer` and  input array
>>> @sk.tree_state.def_state(LayerWithState)
... def _(leaf, *, input: jax.Array) -> jax.Array:
...    return jax.random.normal(jax.random.key(0), input.shape)
>>> sk.tree_state(LayerWithState(), input=jax.numpy.ones((1, 1))).shape
(1, 1)

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> tree = [1, 2, sk.nn.BatchNorm(5, key=jr.key(0))]
>>> sk.tree_state(tree)
[NoState(), NoState(), BatchNormState(
  running_mean=f32[5](μ=0.00, σ=0.00, ∈[0.00,0.00]),
  running_var=f32[5](μ=1.00, σ=0.00, ∈[1.00,1.00])
)]

serket.tree_eval(tree)

Modify tree layers to disable any trainning related behavior.

For example, nn.Dropout layer is replaced by an nn.Identity layer and nn.BatchNorm layer is replaced by EvalBatchNorm layer when evaluating the tree.

tree_eval() objective is to provide a simple and consistent way to disable any trainning related behavior for a tree of layers. Specifically, it provides a way to separate the evaluation logic from the layer definition. This allows for a more clear separation of concerns, and makes it easier to define new layers that has a single behavior.

Parameters:

tree – A tree of layers.

Returns:

A tree of layers with evaluation behavior of same structure as tree.

Example

>>> # dropout is replaced by an identity layer in evaluation mode
>>> # by registering `tree_eval.def_eval(sk.nn.Dropout, sk.nn.Identity)`
>>> import jax.numpy as jnp
>>> import serket as sk
>>> layer = sk.nn.Dropout(0.5)
>>> sk.tree_eval(layer)
Identity()

Note

To define evaluation rule for a custom layer, use the decorator tree_eval.def_eval() on a function that accepts the layer. The function should return the evaluation layer.

>>> import serket as sk
>>> import jax
>>> class AddOne(sk.TreeClass):
...    def __call__(self, input: jax.Array) -> jax.Array:
...        return input + 1
>>> input = jax.numpy.ones([3, 3])
>>> add_one = AddOne()
>>> print(add_one(input))  # add one to each element
[[2. 2. 2.]
 [2. 2. 2.]
 [2. 2. 2.]]

>>> class AddOneEval(sk.TreeClass):
...    def __call__(self, input: jax.Array) -> jax.Array:
...        return input  # no-op

>>> # register `AddOne` to be replaced by `AddOneEval` in evaluation mode
>>> @sk.tree_eval.def_eval(AddOne)
... def _(_: AddOne) -> AddOneEval:
...    return AddOneEval()
>>> print(sk.tree_eval(add_one)(input))
[[1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]]