class FFNN(nn.Module):
n_layers: int
width: int
dtype: Any = np.float64
activation: Any = jax.nn.selu
kernel_init: NNInitFunc = variance_scaling(1.0, "fan_in", "normal")
bias_init: NNInitFunc = normal(0.01)
def setup(self):
self.layers = [
nk.nn.Dense(
features=self.width,
use_bias=True,
dtype=self.dtype,
kernel_init=self.kernel_init,
bias_init=self.bias_init,
) for layer in range(self.n_layers)
]
@nn.compact
def __call__(self, x_in):
x = x_in
for layer in self.layers:
x = layer(x)
x = self.activation(x)
return jnp.sum(x, axis=-1) / (x.shape[-1])**0.5
kernel,
strides,
self.padding,
rhs_dilation=self.kernel_dilation,
precision=self.precision)
if is_single_input:
y = jnp.squeeze(y, axis=0)
if self.use_bias:
bias = self.param('bias', self.bias_init, (self.features,))
bias = jnp.asarray(bias, self.dtype)
y = y + bias
return y
default_embed_init = variance_scaling(1.0, 'fan_in', 'normal', out_axis=0)
class Embed(Module):
"""Embedding Module.
A parameterized function from integers [0, n) to d-dimensional vectors.
Attributes:
num_embeddings: number of embeddings.
features: number of feature dimensions for each embedding.
dtype: the dtype of the embedding vectors (default: float32).
embedding_init: embedding initializer.
"""
num_embeddings: int
features: int