RNN Compiler Walkthrough

RNN Compiler Walkthrough#

This tutorial follows the same compiler workflow from a single recurrent layer to a stacked RNN. We first establish how to read the compilation report and ETraceGraph, then compile a two-layer model and compare the structures reported for both models.

Single-Layer RNN#

Define and Compile the Model#

The model contains one ValinaRNNCell and one linear readout. The recurrent cell owns the temporal hidden state, while the readout maps that state to the final output.

import jax
import jax.numpy as jnp
import brainstate
import braintrace

class SingleLayerRNN(brainstate.nn.Module):
    def __init__(self, n_in, n_rec, n_out):
        super().__init__()
        self.rnn = braintrace.nn.ValinaRNNCell(n_in, n_rec)
        self.out = braintrace.nn.Linear(n_rec, n_out)

    def update(self, x):
        return self.out(self.rnn(x))


model = SingleLayerRNN(10, 32, 5)

# braintrace.compile initialises states, compiles the ETP graph, and returns a ready learner.
# We compile for a single unbatched sample (no batch_size), so the hidden state is (32,) and
# the recurrent op is the matrix-vector primitive etp_mv. verbose=2 prints full diagnostics.
learner = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros(10), verbose=2)
learner.show_graph()

========================================================================================================================
The hidden groups are:

   Group 0: [('rnn', 'h')]


The weight parameters which are associated with the hidden states are:

   Weight 0: ('rnn', 'W', 'weight')  is associated with hidden group 0


The non-etrace weight parameters are:

   Weight 0: ('out', 'weight')  (excluded: relation_excluded_non_temporal)


Compiler diagnostics (warnings / errors):

   [warning] relation_excluded_non_temporal: ETP primitive etp_mv (weight=('out', 'weight')) has no connected hidden states. It will be treated as a non-temporal parameter.



========================================================================================================================
The hidden groups are:

   Group 0: [('rnn', 'h')]


The weight parameters which are associated with the hidden states are:

   Weight 0: ('rnn', 'W', 'weight')  is associated with hidden group 0


The non-etrace weight parameters are:

   Weight 0: ('out', 'weight')  (excluded: relation_excluded_non_temporal)

Read the Compiler Output#

The compiler identifies one hidden group at ('rnn', 'h'). The recurrent weight at ('rnn', 'W', 'weight') reaches that hidden state through an ETP primitive, so the learner maintains an eligibility trace for it.

The readout weight at ('out', 'weight') is different. Its output does not feed a recurrent hidden state, so it is reported as relation_excluded_non_temporal. It remains trainable through the loss, but it does not require a temporal eligibility trace.

Using learner.report – the CompilationReport#

Every learner returned by braintrace.compile exposes a CompilationReport through learner.report. The report is the concise view of what compilation included, excluded, or diagnosed:

  • report.counts summarizes hidden groups, ETP weights, excluded weights, warnings, and errors.

  • report.etrace_weights lists parameter paths that participate in eligibility tracing.

  • report.excluded_weights pairs excluded parameter paths with their reasons.

  • report.dynamic_states records non-hidden dynamic states discovered during tracing.

  • report.diagnostics contains the complete CompilationRecord sequence.

# report.show(level) prints a structured summary at the requested verbosity.
# level=1 shows hidden groups, etrace weights, and excluded weights.
learner.report.show(1)

# Programmatic access to the summary counts
print("Counts:", learner.report.counts)

# Which weights participate in online learning?
print("ETrace weights:", learner.report.etrace_weights)

# Which weights were excluded (e.g., non-temporal readouts)?
print("Excluded weights:", learner.report.excluded_weights)

========================================================================================================================
The hidden groups are:

   Group 0: [('rnn', 'h')]


The weight parameters which are associated with the hidden states are:

   Weight 0: ('rnn', 'W', 'weight')  is associated with hidden group 0


The non-etrace weight parameters are:

   Weight 0: ('out', 'weight')  (excluded: relation_excluded_non_temporal)



Counts: {'hidden_groups': 1, 'etrace_weights': 2, 'excluded_weights': 1, 'warnings': 1, 'errors': 0}
ETrace weights: [(('rnn', 'W', 'weight'), [0]), (('rnn', 'W', 'weight'), [0])]
Excluded weights: [(('out', 'weight'), 'relation_excluded_non_temporal')]

Understanding ETraceGraph#

The report summarizes decisions; learner.graph exposes their structural representation. Its central fields are:

Field

Meaning

module_info

Traced Jaxpr and model-state mappings

hidden_groups

Hidden states updated as recurrent groups

hid_path_to_group

Hidden-state path to group mapping

hidden_param_op_relations

Parameter, ETP primitive, and hidden-group relations

hidden_perturb

Perturbation structure used for hidden Jacobians

diagnostics

Included and excluded compiler decisions

Inspecting the relation itself confirms which trainable path, primitive, and hidden group produced the summary above.

graph = learner.graph

print("=== Hidden Groups ===")
for g in graph.hidden_groups:
    print(f"  Group {g.index}: {g.num_state} state(s), shape {g.varshape}")
    print(f"    Paths: {g.hidden_paths}")

print("\n=== Weight-Primitive-Hidden Relations ===")
for i, r in enumerate(graph.hidden_param_op_relations):
    print(f"  Relation {i}:")
    # ``trainable_paths`` is a dict {trainable key -> owning ParamState path};
    # a single primitive may own several (e.g. {weight, bias}).
    print(f"    Trainable paths: {r.trainable_paths}")
    print(f"    Primitive: {r.primitive}")
    print(f"    Hidden groups: {[g.index for g in r.hidden_groups]}")

print(f"\n=== Perturbation ===")
print(f"  Has perturbation: {graph.hidden_perturb is not None}")

=== Hidden Groups ===
  Group 0: 1 state(s), shape (32,)
    Paths: [('rnn', 'h')]

=== Weight-Primitive-Hidden Relations ===
  Relation 0:
    Trainable paths: {'weight': ('rnn', 'W', 'weight'), 'bias': ('rnn', 'W', 'weight')}
    Primitive: etp_mv
    Hidden groups: [0]

=== Perturbation ===
  Has perturbation: True

Using compile_etrace_graph Directly#

braintrace.compile(...) is the standard entry point because it initializes states, compiles the graph, and returns a ready learner. For structural debugging or custom algorithm development, braintrace.compile_etrace_graph(...) returns the same ETraceGraph without constructing an algorithm wrapper.

model_direct = SingleLayerRNN(10, 32, 5)
brainstate.nn.init_all_states(model_direct)

graph_direct = braintrace.compile_etrace_graph(model_direct, jnp.zeros(10))

print(f"Number of hidden groups: {len(graph_direct.hidden_groups)}")
print(f"Number of relations: {len(graph_direct.hidden_param_op_relations)}")
print(f"Has perturbation: {graph_direct.hidden_perturb is not None}")

print("\nGraph fields:")
for key in graph_direct.dict().keys():
    print(f"  {key}")

Number of hidden groups: 1
Number of relations: 1
Has perturbation: True

Graph fields:
  module_info
  hidden_groups
  hid_path_to_group
  hidden_param_op_relations
  hidden_perturb
  diagnostics

Two-Layer RNN#

The stacked model uses the same input and hidden width but introduces a second recurrent state. It uses GRUCell layers to expose multiple gate relations as well as multiple hidden groups; consequently, the final comparison is a comparison of the compiled structures, not a controlled claim that every difference is caused only by depth.

class TwoLayerRNN(brainstate.nn.Module):
    def __init__(self, n_in, n_rec, n_out):
        super().__init__()
        self.rnn1 = braintrace.nn.GRUCell(n_in, n_rec)
        self.rnn2 = braintrace.nn.GRUCell(n_rec, n_rec)
        self.out = braintrace.nn.Linear(n_rec, n_out)

    def update(self, x):
        h1 = self.rnn1(x)
        h2 = self.rnn2(h1)
        return self.out(h2)


model2 = TwoLayerRNN(10, 32, 5)
learner2 = braintrace.compile(model2, braintrace.D_RTRL, jnp.zeros(10))
learner2.show_graph()

Inspect the Stacked Graph#

The compiler should keep the two sequential recurrent states in separate hidden groups. For each GRU layer, the update-gate (Wz) and candidate (Wh) weights have direct ETP relations to that layer’s hidden group. Reset-gate (Wr) paths are excluded when they reach the hidden state through another trainable ETP primitive, while the readout remains non-temporal.

graph2 = learner2.graph

print(f"Number of hidden groups: {len(graph2.hidden_groups)}")
print(f"Number of weight-hidden relations: {len(graph2.hidden_param_op_relations)}")
print("Hidden paths:", [group.hidden_paths for group in graph2.hidden_groups])
print("Report counts:", learner2.report.counts)

Single-Layer vs Two-Layer#

The comparison below is generated from the two compiled learners rather than maintained as a separate hand-written result table. It distinguishes model architecture from the compiler structures discovered for each model.

single_counts = learner.report.counts
two_counts = learner2.report.counts

print(f"{'Property':<32} {'Single layer':>14} {'Two layers':>14}")
print("-" * 62)
print(f"{'Recurrent cell':<32} {'ValinaRNNCell':>14} {'GRUCell':>14}")
print(f"{'Hidden groups':<32} {len(learner.graph.hidden_groups):>14} {len(learner2.graph.hidden_groups):>14}")
print(f"{'Weight-hidden relations':<32} {len(learner.graph.hidden_param_op_relations):>14} {len(learner2.graph.hidden_param_op_relations):>14}")
print(f"{'ETP weights':<32} {single_counts['etrace_weights']:>14} {two_counts['etrace_weights']:>14}")
print(f"{'Excluded weights':<32} {single_counts['excluded_weights']:>14} {two_counts['excluded_weights']:>14}")

The second recurrent layer creates another independently represented hidden-state transition. The additional GRU gates also create more candidate parameter paths, so the relation-count difference must not be attributed to depth alone. In both models, the report explains classification outcomes and the graph identifies the exact parameter, primitive, and hidden-group connections behind them.

Summary#

The single-layer model establishes the inspection sequence; the two-layer model applies it without repeating the API explanation. Comparing the resulting reports and graphs shows why compiler decisions must be interpreted at the relation level rather than inferred from module names or parameter counts. These structural checks establish what the compiler selected, but they do not by themselves establish gradient correctness.