Linear

class serket.nn.Linear(in_features, out_features, *, key, in_axis=-1, out_axis=-1, weight_init='glorot_uniform', bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

Apply a Linear Layer to input.

Parameters:

• in_features (Union[int, Sequence[int], None]) – number of input features corresponding to in_axis. Accepts int or tuple of ints.

• out_features (Union[int, Sequence[int]]) – number of output features corresponding to out_axis. Accepts int or tuple of ints.

• key (Array) – key to use for initializing the weight and biases.

• in_axis (Union[int, Sequence[int]]) – axes to apply the linear layer to. Accepts int or tuple of ints. Defaults to -1.

• out_axis (Union[int, Sequence[int]]) – result axis. Accepts int or tuple of ints. Defaults to -1.

• weight_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – weight initialization function. Defaults to glorot_uniform.

• bias_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – bias initialization function. Defaults to zeros.

• dtype (Union[dtype, str, Any]) – dtype of the weight and biases. float32

Example

Apply Linear layer t0 the last dimension of input

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> input = jnp.ones([1, 2, 3, 4])
>>> key = jr.key(0)
>>> layer = sk.nn.Linear(4, 5, key=key)
>>> layer(input).shape
(1, 2, 3, 5)

Example

Apply Linear layer to first and second axes of input

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> input = jnp.ones([1, 2, 3, 4])
>>> in_axis = (0, 1)  # which axes to apply the linear layer to
>>> in_features = (1, 2)  # number of input features corresponding to ``in_axis``
>>> out_axis = (0, 2) # which axes to map the output to
>>> out_features = (3, 4)  # number of output features corresponding to ``out_axis``
>>> key = jr.key(0)
>>> layer = sk.nn.Linear(in_features, out_features, in_axis=in_axis, out_axis=out_axis, key=key)
>>> layer(input).shape
(3, 3, 4, 4)

Note

Linear supports lazy initialization, meaning that the weight and biases are not initialized until the first call to the layer. This is useful when the input shape is not known at initialization time.

To use lazy initialization, pass None as the in_features argument and use value_and_tree() to call the layer and return the method output and the material layer.

>>> import jax
>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> input = jnp.ones((10, 5, 4))
>>> lazy = sk.nn.Linear(None, 12, in_axis=(0, 2), key=key)
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> material.in_features
(10, 4)

__call__(input)

Apply a linear transformation to the input.

Parameters:

input (Array)

Return type:

Array

static linear_op(weight, bias, in_axis=(-1,), out_axis=(-1,))

Apply a linear transformation to the input.

Parameters:

• input (Array) – input array.

• weight (Array) – weight array with shape (out_feature_1, out_feature_2, …, in_feature_1, in_feature_2, …). in_feature_i corresponds to in_axis[i] and out_feature_i corresponds to out_axis[i].

• bias (Array | None) – bias array with shape (out_feature_1, out_feature_2, …) or None for no bias.

• in_axis (Sequence[int]) – axes to apply the linear layer to.

• out_axis (Sequence[int]) – result axis.

Return type:

Array

class serket.nn.Identity(*a, **k)

Identity layer. Returns the input.

__call__(input, **_)

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.Embedding(in_features, out_features, key)

Defines an embedding layer.

Parameters:

• in_features (int) – vocabulary size.

• out_features (int) – embedding size.

• key (Array) – random key to initialize the weight.

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> # 10 words in the vocabulary, each word is represented by a 3 dimensional vector
>>> key = jr.key(0)
>>> table = sk.nn.Embedding(10, 3, key=key)
>>> # take the last word in the vocab
>>> input = jnp.array([9])
>>> output = table(input)
>>> output.shape
(1, 3)

__call__(input)

Embeds the input.

Parameters:

input (Array) – integer index input of subdtype integer.

Return type:

Array

Returns:

Embedding of the input.

class serket.nn.MLP(in_features, out_features, hidden_features, num_hidden_layers, *, key, act='tanh', weight_init='glorot_uniform', bias_init='zeros', dtype=<class 'jax.numpy.float32'>)

Multi-layer perceptron.

Parameters:

• in_features (int) – Number of input features.

• out_features (int) – Number of output features.

• hidden_features (int) – Number of hidden units in each hidden layer.

• num_hidden_layers (int) – Number of hidden layers including the output layer.

• key (Array) – Random number generator key.

• act (Union[Literal['celu', 'elu', 'gelu', 'glu', 'hard_shrink', 'hard_sigmoid', 'hard_swish', 'hard_tanh', 'leaky_relu', 'log_sigmoid', 'log_softmax', 'mish', 'prelu', 'relu', 'relu6', 'selu', 'sigmoid', 'softplus', 'softshrink', 'softsign', 'squareplus', 'swish', 'tanh', 'tanh_shrink', 'thresholded_relu'], Callable[[Union[Array, ndarray, bool, number, bool, int, float, complex]], Array]]) – Activation function. Defaults to tanh.

• weight_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – Weight initialization function. Defaults to glorot_uniform.

• bias_init (Union[Literal['he_normal', 'he_uniform', 'glorot_normal', 'glorot_uniform', 'lecun_normal', 'lecun_uniform', 'normal', 'uniform', 'ones', 'zeros', 'xavier_normal', 'xavier_uniform', 'orthogonal'], Callable[[Array, Tuple[int, ...], Union[dtype, str, Any]], Array | None]]) – Bias initialization function. Defaults to zeros.

• dtype (Union[dtype, str, Any]) – dtype of the weight and biases. float32

Example

>>> import jax.numpy as jnp
>>> import serket as sk
>>> import jax.random as jr
>>> key = jr.key(0)
>>> layer = sk.nn.MLP(1, 2, hidden_features=4, num_hidden_layers=2, key=key)
>>> input = jnp.ones((3, 1))
>>> layer(input).shape
(3, 2)

Note

• MLP with in_features=1, out_features=2, hidden_features=4, num_hidden_layers=2 is equivalent to [1, 4, 4, 2] which has one input layer (1, 4), one intermediate layer (4, 4), and one output layer (4, 2) = num_hidden_layers + 1

Note

MLP supports lazy initialization, meaning that the weight and biases are not initialized until the first call to the layer. This is useful when the input shape is not known at initialization time.

To use lazy initialization, pass None as the in_features argument and use value_and_tree() to call the layer and return the method output and the material layer.

>>> import serket as sk
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> key = jr.key(0)
>>> lazy = sk.nn.MLP(None, 1, num_hidden_layers=2, hidden_features=10, key=key)
>>> input = jnp.ones([1, 10])
>>> _, material = sk.value_and_tree(lambda lazy: lazy(input))(lazy)
>>> material.in_features
10

Note

MLP uses jax.lax.scan to reduce the jaxpr size. Leading to faster compilation times and smaller jaxpr size.

>>> import serket as sk
>>> import jax
>>> import jax.numpy as jnp
>>> # 10 hidden layers
>>> mlp1 = sk.nn.MLP(1, 2, 5, 10, key=jax.random.key(0))
>>> # 50 hidden layers
>>> mlp2 = sk.nn.MLP(1, 2, 5, 50, key=jax.random.key(0))
>>> jaxpr1 = jax.make_jaxpr(mlp1)(jnp.ones([10, 1]))
>>> jaxpr2 = jax.make_jaxpr(mlp2)(jnp.ones([10, 1]))
>>> # same number of equations irrespective of the number of hidden layers
>>> assert len(jaxpr1.jaxpr.eqns) == len(jaxpr2.jaxpr.eqns)

__call__(input)

Call self as a function.

Parameters:

input (Array)

Return type:

Array

serket.nn.linear(input, weight, bias, in_axis=(-1,), out_axis=(-1,))

Apply a linear transformation to the input.

Parameters:

• input (Array) – input array.

• weight (Array) – weight array with shape (out_feature_1, out_feature_2, …, in_feature_1, in_feature_2, …). in_feature_i corresponds to in_axis[i] and out_feature_i corresponds to out_axis[i].

• bias (Array | None) – bias array with shape (out_feature_1, out_feature_2, …) or None for no bias.

• in_axis (Sequence[int]) – axes to apply the linear layer to.

• out_axis (Sequence[int]) – result axis.

Return type:

Array