gensbi.models.simformer.transformer#

Classes#

`AttentionBlock`	Self-attention block with a fixed fp32 compute island.
`DenseBlock`	MLP block (the FLOPs-dominant part of the transformer) with an fp32
`Transformer`	A transformer stack.

Module Contents#

class gensbi.models.simformer.transformer.AttentionBlock(din, num_heads, features, skip_connection, rngs, dtype=jnp.float32, param_dtype=jnp.float32)[source]#

Bases: flax.nnx.Module

Self-attention block with a fixed fp32 compute island.

flax.nnx.MultiHeadAttention (flax 0.12.7) computes its softmax in whatever dtype it is given – there is no internal fp32 upcast for the attention logits/softmax the way some other implementations provide. Rather than bolt on a custom attention_fn just to force fp32 softmax, this block keeps all of its internal math (LayerNorm + MultiHeadAttention) at dtype=jnp.float32, ignoring the dtype compute-precision knob for those internals. This is a deliberately larger fp32 island than a single softmax op, but the bulk of the model’s FLOPs live in the DenseBlock MLP stack (widening_factor x wider), which does honor the bf16 dtype knob, so this island has a small cost in practice.

The island is for the math only: the block’s output (after the optional skip connection) is downcast to the requested dtype before being returned, mirroring the codebase’s established fp32-island idiom (e.g. QKNorm.__call__’s .astype(v.dtype) in flux1/layers.py). Without this downcast, the fp32 residual (x_in, captured post fp32-LayerNorm) would silently re-promote every downstream block’s output back to fp32 via JAX’s bf16+fp32 promotion rule on the skip-add, defeating the bf16 knob’s memory/bandwidth benefit for the whole inter-block residual stream.

param_dtype (master-weight storage) is unaffected and still threads through normally.

Parameters:

din (int)
num_heads (int)
features (int)
skip_connection (bool)
rngs (flax.nnx.Rngs)
dtype (jax.typing.DTypeLike)
param_dtype (jax.typing.DTypeLike)

__call__(x, mask)[source]#

Parameters:

x (jax.numpy.ndarray)
mask (jax.numpy.ndarray | None)

Return type:

jax.numpy.ndarray

attn[source]#

dtype[source]#

layer_norm[source]#

skip_connection[source]#

class gensbi.models.simformer.transformer.DenseBlock(din, dcontext, num_hidden_layers, widening_factor, act, skip_connection, rngs, dtype=jnp.float32, param_dtype=jnp.float32)[source]#

Bases: flax.nnx.Module

MLP block (the FLOPs-dominant part of the transformer) with an fp32 LayerNorm island.

The LayerNorm math runs in fp32 (mirrors AttentionBlock’s normalization treatment); the wide hidden-layer matmuls run in the requested compute dtype. The block’s output (after the context-merge and optional skip connection) is downcast to dtype before being returned – otherwise the fp32 residual (x_in, captured post fp32-LayerNorm) would silently re-promote the output back to fp32 via JAX’s bf16+fp32 promotion rule on the skip-add, defeating the bf16 knob’s memory/bandwidth benefit for the whole inter-block residual stream (see AttentionBlock docstring for the same pattern).

Parameters:

widening_factor (int)
act (Callable)
skip_connection (bool)
rngs (flax.nnx.Rngs)
dtype (jax.typing.DTypeLike)
param_dtype (jax.typing.DTypeLike)

__call__(x, context)[source]#

act[source]#

context_block[source]#

dtype[source]#

hidden_blocks[source]#

layer_norm[source]#

skip_connection[source]#

class gensbi.models.simformer.transformer.Transformer(din, dcontext, num_heads, num_layers, features, widening_factor=4, num_hidden_layers=1, act=jax.nn.gelu, skip_connection_attn=True, skip_connection_mlp=True, *, rngs, dtype=jnp.float32, param_dtype=jnp.float32)[source]#