Gradient Optimization Circuit Search Jacobians Stream Calculus Rewriting Multi-Dim Streams Coupled SDEs JAX & JIT Engineering Examples

Stream Calculus Internals

How Asgard solves differential equations using Rutten's coinductive stream calculus.

Key Idea

A stream is an infinite sequence of coefficients:

f = [c₀, c₁, c₂, c₃, ...]

When interpreted as Taylor coefficients, this represents the function f(x) = c₀ + c₁x + c₂x²/2! + .... Asgard solves differential equations by computing these coefficients one at a time, using feedback to handle recursion.

Register and Deregister

These are the fundamental shift operations on streams.

Register (integration)

Shifts the stream right and prepends zero:

register: [c₀, c₁, c₂, ...] → [0, c₀, c₁, ...]

This corresponds to integration: if f has Taylor coefficients [c₀, c₁, ...] then ∫f has coefficients [0, c₀, c₁, ...] (with appropriate scaling).

Deregister (differentiation)

Shifts the stream left, dropping the first element:

deregister: [c₀, c₁, c₂, ...] → [c₁, c₂, c₃, ...]

This corresponds to differentiation: the first coefficient is discarded because the constant term vanishes under d/dx.

Inverse Relationship

Register and deregister are inverses (up to the lost constant):

deregister(register(f)) = f
register(deregister(f)) = f - c₀   (constant term lost)

Trace (Feedback)

The trace operator creates a feedback loop, turning a differential equation into a fixed-point problem that can be solved coefficient by coefficient.

How It Works

For a circuit f: (n+1) → (m+1), trace(f): n → m feeds the last output wire back to the last input wire:

         ┌──────────────────────┐
         │                      │
  n ───▶ │         f            │ ───▶ m
         │                      │
         │   ┌──────────────┐   │
         │   │  feedback     │◀──┘
         │   └──────────────┘
         └──────────────────────┘

Example: Exponential Function

The equation f' = f with f(0) = 1 becomes f = const(1) + register(f):

Circuit: trace(composition(monoidal(const(1.0), register(x)), add))

The evaluator computes coefficients step by step:

Step 0: c₀ = const(1)[0] + register(feedback)[0]
           = 1          + 0             = 1
Step 1: c₁ = const(1)[1] + register(feedback)[1]
           = 0          + feedback[0]   = 1
Step 2: c₂ = const(1)[2] + register(feedback)[2]
           = 0          + feedback[1]   = 1
...

Result: [1, 1, 1, 1, ...] — the Taylor coefficients of eˣ (before factorial scaling).

Coefficient-by-Coefficient Evaluation

Why Not Just Use Arrays?

Standard array-based computation processes all elements at once. This fails for feedback loops: you can't compute position i of the output until you know position i-1 of the feedback, which itself depends on all earlier positions.

The StreamCalculusEvaluator solves this by building the result one coefficient at a time:

  1. Start with empty output buffers
  2. For each index i from 0 to n:
    • Compute the circuit output at position i using the prefix of already-computed feedback coefficients
    • Store the result in the output buffer
  3. Return the complete coefficient arrays

Multiplication as Cauchy Product

When two streams are multiplied, the coefficient at position i is the Cauchy product (discrete convolution):

(f · g)[i] = Σ(k=0 to i) f[k] · g[i-k]

This naturally fits coefficient-by-coefficient evaluation since (f · g)[i] only depends on coefficients up to position i.

When Stream Calculus Is Used

The JAXCircuitCompiler dispatches to StreamCalculusEvaluator when a circuit has both a trace operator and a register operator:

if circuit.has_trace and circuit.has_register:
    # → StreamCalculusEvaluator (coefficient-by-coefficient)
else:
    # → Standard JAX compilation (array-based)

This combination indicates a differential equation (feedback + integration). Standard JAX fixed-point iteration would converge incorrectly or not at all for these circuits.

Trace Wiring

The compilation pipeline automatically converts implicit recursion into explicit trace operators:

  1. Variable isolation — identifies the unknown function in the equation
  2. Equation-to-circuit — translates to a circuit with var(f) nodes
  3. Trace wiring — replaces duplicate var(f) with id, adds split at the output, and wraps with trace()

For example, f = f₀ + ∫f compiles through these stages:

Equation:  int(f, x) + f0 = f
Circuit:   composition(monoidal(var(f0), composition(var(f), register(x))), add)
Wired:     trace(composition(composition(monoidal(var(f0),
               composition(id, register(x))), add), split))

Limitations