gensbi.models.maf.masked_linear

Dense layer with a fixed binary weight mask.

Classes

MaskedLinear

y = (kernel * mask).T @ x + bias with a non-trainable mask.

Module Contents

class gensbi.models.maf.masked_linear.MaskedLinear(in_features, out_features, mask, rngs, param_dtype=jnp.float32, dtype=jnp.float32)

Bases: flax.nnx.Module

y = (kernel * mask).T @ x + bias with a non-trainable mask.

Parameters:

• in_features (int)

• out_features (int)

• mask (Array) – Boolean array of shape (in_features, out_features); stored as a Mask buffer so it is excluded from nnx.Param.

• rngs (nnx.Rngs)

• param_dtype (DTypeLike, optional) – Dtype for the stored (master) kernel/bias parameters. Defaults to float32 (exact-likelihood model needs the precision).

• dtype (DTypeLike, optional) – Compute dtype: the kernel/bias/activations are cast to this dtype before the matmul. Defaults to float32, matching param_dtype, so with default arguments this is a no-op cast (bit-identical).

__call__(x)

Apply the masked linear transform y = (kernel * mask).T @ x + bias.

Parameters:

x (Array) – Input of shape (in_features,).

Returns:

Output of shape (out_features,).

Return type:

Array

dtype

linear

mask