dockyard_rl.models.jax.layers

Reusable Flax NNX primitives built on explicit nnx.Param arrays.

Each primitive stores its weight as a raw nnx.Param in JAX-native layout (linear kernels as (in, out)), so the HF-weights loader name-map, the J2 sharding PartitionSpec attachment, and the J5 refit inverse are all explicit and live in one place rather than threaded through framework layer internals. The per-projection boilerplate is written once here and composed everywhere.

A sharding metadata tuple may be attached to each nnx.Param (logical mesh-axis names per array dim, or None for replicated). It is unused at J1 (single-device parity) and consumed by J2.

Module Contents

Classes

Linear	y = x @ kernel with kernel stored (in_features, out_features).
Embedding	Token embedding table (num_embeddings, features); gather by id.
RMSNorm	RMSNorm matching Qwen3RMSNorm numerics exactly.

Data

Array
Sharding

API

dockyard_rl.models.jax.layers.Array

None

dockyard_rl.models.jax.layers.Sharding

None

class dockyard_rl.models.jax.layers.Linear(in_features: int, out_features: int, *, rngs: flax.nnx.Rngs, param_dtype: jax.numpy.dtype = jnp.float32, kernel_init: flax.typing.Initializer = nnx.initializers.lecun_normal(), sharding: dockyard_rl.models.jax.layers.Sharding = None)

Bases: flax.nnx.Module

y = x @ kernel with kernel stored (in_features, out_features).

Bias-free: every Qwen3 dense projection (q/k/v/o, gate/up/down, lm_head) has bias=False. HF nn.Linear stores weight as (out, in); the loader transposes into this (in, out) convention.

Initialization

class dockyard_rl.models.jax.layers.Embedding(num_embeddings: int, features: int, *, rngs: flax.nnx.Rngs, param_dtype: jax.numpy.dtype = jnp.float32, embedding_init: flax.typing.Initializer = nnx.initializers.normal(stddev=1.0), sharding: dockyard_rl.models.jax.layers.Sharding = None)

Bases: flax.nnx.Module

Token embedding table (num_embeddings, features); gather by id.

Matches HF nn.Embedding.weight layout exactly (no transpose on load).

Initialization

class dockyard_rl.models.jax.layers.RMSNorm(dim: int, *, eps: float = 1e-06, param_dtype: jax.numpy.dtype = jnp.float32, sharding: dockyard_rl.models.jax.layers.Sharding = None)

Bases: flax.nnx.Module

RMSNorm matching Qwen3RMSNorm numerics exactly.

Variance is computed in float32, the normalized activation is cast back to the input dtype, then multiplied by the (param-dtype) scale — identical to HF’s self.weight * hidden_states.to(input_dtype) ordering, which matters for bf16 parity. The scale initializes to ones.

Initialization