nnterp · encyclopedia jamba · nnterp.families.jamba

family jamba

JambaJamba 1.5–1.7Jamba2

Llama's pre-norm block with one of two mixers, a Mamba-1 selective scan with RMS norms on its step, B and C on most blocks and attention without rotary on one block in eight or fourteen, and, on every other block of the larger checkpoints, a mixture of experts in place of the MLP.

HybridSelective scan (Mamba-1)Blocks without rotaryMixture of expertsUnnormalized routingGrouped-query attention

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card. The slider picks a block, each tick's color gives its block's shape, and the diagram redraws for the block picked.

embed_tokenstoken_embeddings
→
layers× 32
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity layers[i].input + linear_attn.attention_output + mlp.mlp_output == layer_output

The API

02 · support()

Values, by host

⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.

model JambaForCausalLMlayoutwhat it isread at
logits Logits [batch seq vocab_size] The logits the model returns, after anything it does past lm_head (a softcap, a scale)

the module's own output
token_embeddings Residual [batch seq hidden_size] The embedding module's output; layers[0].input is what enters the first block

the child module embed_tokens's output
final_norm_output Residual [batch seq hidden_size] The residual stream after the final norm, what lm_head reads

the child module norm's output
lm_head_output Logits [batch seq vocab_size] The head's logits before anything the model does past lm_head (a softcap, a scale)

the child module lm_head's output
next_token_probs NextTokenProbs [batch vocab_size] The next-token distribution at the last position; derived, read-only

the module's own output
input_ids Tokens [batch seq] The token ids the model was called with

the module's own input
attention_mask Tokens [batch seq] The attention mask the model was called with; zeros are padding

the module's own input
input_size scalar [batch, seq] of the current call; read-only

the module's own input
model.layers[i] JambaMambaDecoderLayerlayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
attention_norm_output Residual [batch seq hidden_size] The normed stream entering the attention

the child module input_layernorm's output
layer_mid Residual [batch seq hidden_size] The residual stream after the attention add

the child module post_attention_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module post_attention_layernorm's output
mlp_input Residual [batch seq hidden_size] What the MLP is called with

the child module mlp's input
layer_output Residual [batch seq hidden_size] The residual stream leaving the block, a tensor even when the block returns a tuple

the module's own output
model.layers[i].linear_attn JambaMambaMixerlayoutwhat it isread at
attention_output Residual [batch seq hidden_size] What the mixer adds to the residual stream

the module's own output
⚠︎state ScanState [batch channels state_dim] The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')

inside the forward, at the operation that _token_state_op finds on each call
⚠︎states ScanStates [batch seq channels state_dim] The state after every token of this call; needs route_kernels(family, 'torch')

inside the forward, at the operation that _scan_states finds on each call
⚠︎attention_queries ScanQK [batch seq groups state_dim] C, the vector each token reads the state with

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), the element that call returns it in
⚠︎attention_keys ScanQK [batch seq groups state_dim] B, the vector each token writes into the state with

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), the element that call returns it in
⚠︎attention_values ScanValues [batch seq channels] x, the input the scan writes into the state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), the element that call returns it in
⚠︎betas ScanSteps [batch seq channels] The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only

inside the forward, at the operation that _betas finds on each call
⚠︎decays ScanDecays [batch seq channels state_dim] The per-token log decay dt * A of the state; derived, read-only

inside the forward, at the operation that _decays finds on each call
⚠︎state_input ScanState | None [batch channels state_dim] The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), the element that call returns it in
⚠︎attention_head_outputs ScanValues [batch seq channels] y = C.h + D.x, the scan's read of the state before the gate and out_proj

inside the forward, at the operation that inside.READ_OPS finds on each call
⚠︎state_output ScanState [batch channels state_dim] The state leaving the layer

inside the forward, at the operation that _state_output_key finds on each call
model.layers[i].self_attn JambaAttentionlayoutwhat it isread at
⚠︎attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of attention_interface_1, element 1
⚠︎attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of attention_interface_1, element 2
⚠︎attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of attention_interface_1, element 3
⚠︎attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward: the input of nn_functional_softmax_0 inside attention_interface_1
attention_premix HeadOutputs [batch seq num_heads head_dim] Every head's output side by side, as the output projection takes it

the child module o_proj's input
attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream

the module's own output
⚠︎attention_probabilities Pattern [batch num_heads query key] The attention pattern the values are mixed with

inside the forward: the output of nn_functional_dropout_0 inside attention_interface_1
⚠︎attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of attention_interface_1, element 0
model.layers[i].mlp JambaMLP / JambaSparseMoeBlocklayoutwhat it isread at
⚠︎mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn's output
⚠︎mlp_neurons Neurons [batch seq intermediate_size] What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise

the child module down_proj's input
mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream

the module's own output
router_logits RouterLogits [batch seq num_experts] The router's logits, one per expert, before the scoring

the child module router's output
expert_weights ExpertWeights [batch seq top_k] The weight each selected expert's output is scaled by

the child module experts's inputs, element 2
expert_indices ExpertIndices [batch seq top_k] The expert each routing slot sends the token to

the child module experts's inputs, element 1
expert_outputs ExpertOutputs [batch seq top_k hidden_size] Each routing slot's weighted expert output

inside experts's forward: the output of weighted_out_view_0 inside experts_forward_1
routed_output Residual [batch seq hidden_size] The routed experts' combined output

the child module experts's output

Printout

What print(model) shows on ai21labs/Jamba-v0.1: the native modules, standard names and every standard value with its layout.

JambaForCausalLM(
  (lm_head): Linear(in_features=4096, out_features=65536, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(65536, 4096, padding_idx=0)
  (layers): ModuleList(
    (0): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (1): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (2): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (3): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (4): JambaAttentionDecoderLayer(
      (self_attn): JambaAttention(
        (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (5): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (6): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (7): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (8): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (9): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (10): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (11): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (12): JambaAttentionDecoderLayer(
      (self_attn): JambaAttention(
        (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (13): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (14): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (15): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (16): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (17): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (18): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (19): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (20): JambaAttentionDecoderLayer(
      (self_attn): JambaAttention(
        (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (21): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (22): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (23): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (24): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (25): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (26): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (27): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (28): JambaAttentionDecoderLayer(
      (self_attn): JambaAttention(
        (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (29): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (30): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (31): JambaMambaDecoderLayer(
      (linear_attn/mamba): JambaMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): JambaRMSNorm((256,), eps=1e-06)
        (b_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (c_layernorm): JambaRMSNorm((16,), eps=1e-06)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (mlp/feed_forward): JambaSparseMoeBlock(
        (router): Linear(in_features=4096, out_features=16, bias=False)
        (experts): JambaExperts(
          (act_fn): SiLUActivation()
        )
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm/pre_ff_layernorm): JambaRMSNorm((4096,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
  )
  (norm): JambaRMSNorm((4096,), eps=1e-06)
  (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale)
  (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block
  (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads
  (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale)
  (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only
  (input_ids) -> Tokens [batch seq]: The token ids the model was called with
  (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding
  (input_size): [batch, seq] of the current call; read-only
)

Config

model.num_layers32
model.hidden_size4096
model.num_heads32
model.num_kv_heads8
model.head_dim128
model.qk_head_dim128
model.vocab_size65536
model.intermediate_size14336
architecturesJambaForCausalLM
num_hidden_layers32
hidden_size4096
intermediate_size14336
num_attention_heads32
num_key_value_heads8
vocab_size65536
max_position_embeddings262144
hidden_actsilu
rms_norm_eps1e-06
tie_word_embeddingsFalse
dtypebfloat16
num_experts16
num_experts_per_tok2
mamba_d_state16
mamba_expand2
attn_layer_period8
attn_layer_offset4
expert_layer_period2
expert_layer_offset1

The Notes

03 · what to know before you run

The block, in order

h   = x + mixer(input_layernorm(x))     # linear_attn (the Mamba mixer) or self_attn
out = h + ffn(pre_ff_layernorm(h))      # mlp: a SwiGLU MLP or a mixture of experts

Two block classes, JambaMambaDecoderLayer and JambaAttentionDecoderLayer, differ only in the mixer, and both return a tensor. Block i is an attention block where i % attn_layer_period == attn_layer_offset and a Mamba block elsewhere; its feed-forward is a mixture where i % expert_layer_period == expert_layer_offset and num_experts is above 1, and a dense MLP elsewhere. From the configs:

  • Jamba-v0.1 and Jamba2-Mini: 32 blocks, attention period 8 offset 4 (blocks 4, 12, 20, 28), experts period 2 offset 1 (every odd block, 16 experts, top 2). The attention blocks are even, so they always have the dense MLP: the shapes are Mamba + MLP, Mamba + MoE and Attention + MLP.
  • Jamba-Reasoning-3B and Jamba2-3B: 28 blocks, attention period 14 offset 7 (blocks 7 and 21), and num_experts 1, so no block has a mixture.

config.layer_types names the mixer per block (linear_attention, full_attention); config.layers_num_experts the experts per block. Pick blocks outside the trace:

kinds = model.config.layer_types
mamba = [i for i, t in enumerate(kinds) if t == "linear_attention"]
attn = [i for i, t in enumerate(kinds) if t == "full_attention"]          # 4, 12, 20, 28 on v0.1
moe = [i for i, n in enumerate(model.config.layers_num_experts) if n > 1]

The contributions are the sublayers' outputs

mamba is linear_attn, feed_forward is mlp (dense or mixture) and pre_ff_layernorm is post_attention_layernorm. Nothing scales or norms a sublayer's output on its way to the stream, so the identity is the plain sum on every shape, exact in float32:

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()   # attention on the tiny checkpoint; 4 on v0.1
    mlp = model.layers[1].mlp.mlp_output.save()
    out = model.layers[1].layer_output.save()

torch.testing.assert_close(x + attn + mlp, out)

On a Mamba block the middle term is model.layers[i].linear_attn.attention_output, the out_proj output of the mixer, the same [batch, seq, hidden_size] layout.

Loading: route the kernels

With mamba_ssm installed, transformers runs the selective scan in its CUDA kernels: every linear_attn value but attention_output is unavailable, and on a CPU no trace runs at all (RuntimeError: Expected u.is_cuda() to be true), even one that reads only layer_output. nnterp.route_kernels(model.family, "torch") before the first trace binds transformers' pure-torch scan, which is a token loop, so state, states, state_after and set_state_after work without another call:

import nnterp
from nnterp import StandardizedTransformer

nnterp.route_kernels(nnterp.families.jamba, "torch")       # before the first trace
model = StandardizedTransformer("ai21labs/Jamba-v0.1", attn_implementation="eager")

The attention interior on the attention blocks needs attn_implementation="eager".

The Mamba mixer norms its step, B and C

The mixer is Mamba-1's, with three RMS norms with learned weights between x_proj and the scan: dt_layernorm on the low-rank step (mamba_dt_rank, 256 on v0.1 and Jamba2-Mini, 160 on the 3B checkpoints), b_layernorm on B and c_layernorm on C. The values are read at the scan's call, after the norms: attention_queries is c_layernorm's output.

mix = model.layers[0].linear_attn
with model.trace(prompt):
    c = mix.c_layernorm.output.save()          # [batch, seq, state_dim]
with model.trace(prompt):
    C = mix.attention_queries.save()           # [batch, seq, 1, state_dim]

assert torch.equal(c, C[:, :, 0])

The mixer is mamba_expand × hidden_size channels wide (8192 on v0.1, 5120 on the 3B checkpoints) with a state of mamba_d_state 16 per channel: attention_values, betas and attention_head_outputs are [batch, seq, channels], decays [batch, seq, channels, 16], the state [batch, channels, 16] and states [batch, seq, channels, 16]. The root's intermediate_size is the MLP's (14336 on v0.1), not the mixer's width.

The mixture's weights are the softmax's top 2, not renormalized

router is a bare nn.Linear, and router_logits is its output. The mixture takes a softmax over all experts in float32 and keeps the num_experts_per_tok largest; nothing renormalizes them, so on v0.1 and Jamba2-Mini a token's two expert_weights sum to less than one. There is no shared expert: routed_output is mlp_output.

moe = model.layers[1].mlp                      # an odd block: a mixture
with model.trace(prompt):
    logits = moe.router_logits.save()          # [batch, seq, num_experts]
    w = moe.expert_weights.save()              # [batch, seq, top_k]
    idx = moe.expert_indices.save()

top = logits.float().softmax(-1).topk(moe.top_k, dim=-1)
torch.testing.assert_close(top.values.to(w.dtype), w)
assert torch.equal(top.indices, idx)

Attention carries no position

No block has a rotary or position embedding: attention_queries is q_proj's output split into heads, unchanged, and order reaches the attention only through the causal mask and the Mamba blocks before it. On v0.1 and Jamba2-Mini 32 query heads share 8 key/value heads of 128, so an edit to attention_values[:, j] reaches query heads 4j to 4j + 3; the 3B checkpoints have 20 query heads over one key/value head.

attn = model.layers[1].self_attn                # attention on the tiny checkpoint; layers[4] on v0.1
with model.trace(prompt):
    q_raw = attn.source.self_q_proj_0.output.save()
    q = attn.attention_queries.save()          # [batch, num_heads, seq, head_dim]

b, s, _ = q_raw.shape
assert torch.equal(q_raw.view(b, s, -1, model.head_dim).transpose(1, 2), q)

The readout, the embeddings and the tokenizer

logits is lm_head.output, and project_on_vocab on the last block's layer_output equals logits exactly. Nothing is added after embed_tokens, so token_embeddings equals layers[0].input. The 3B checkpoints tie lm_head to embed_tokens; v0.1 and Jamba2-Mini do not. The tokenizer has 65536 tokens and prepends <|startoftext|> (id 1), so position 0 is that token. On v0.1 <|endoftext|> (id 2) ends a text; the Jamba2 and Reasoning-3B tokenizers end a turn with <|im_end|> (id 519), which their configs name as eos_token_id. Jamba-v0.1 is a base model and ships no chat template.

From the family module

04 · nnterp/families/jamba.py

Jamba (JambaForCausalLM): Mamba-1 blocks, attention blocks and a mixture of experts.

Its tree is model.{embed_tokens, layers[i], final_layernorm} plus lm_head, with two block classes (config.layers_block_type): most blocks are JambaMambaDecoderLayer, {input_layernorm, mamba, pre_ff_layernorm, feed_forward}, and one in attn_layer_period is JambaAttentionDecoderLayer, the same with self_attn in place of mamba. Both are Llama's pre-norm sequential block, so the identity is Llama's. Each block has one mixer, never both: on a Mamba block every self_attn value is reported missing and the selective-scan values live at layers[i].linear_attn (see nnterp.SelectiveScan); on an attention block the reverse. feed_forward is a dense MLP or, one block in expert_layer_period, a sparse mixture of experts; both are mlp, and pre_ff_layernorm is post_attention_layernorm.

The attention has no positional encoding (the Mamba blocks carry position) and runs the shared eager interface, so the base holds. The Mamba mixer puts RMS norms on the step size, B and C before the scan (with weights, where Falcon-Mamba's are weightless); the values are read at the kernel call, after them. With mamba_ssm installed, call nnterp.route_kernels(model.family, "torch") before the first trace.