Jacobian Computation

Sensitivity analysis and input-output dependencies using automatic differentiation.

Overview

The Jacobian matrix J[i,j] = ∂output[i]/∂input[j] shows how each output component depends on each input component. This enables:

• Sensitivity analysis - Which inputs matter most?
• Stability analysis - How sensitive is the output to perturbations?
• Feature importance - Which dimensions drive the output?
• Gradient fields - Visualize how outputs change with inputs

Basic Usage

from gimle.asgard.circuit.circuit import Circuit
from gimle.asgard.circuit.circuit_gradients import compute_jacobian
from gimle.asgard.runtime.stream import Stream
import jax.numpy as jnp

# Create a simple scaling circuit
circuit = Circuit.from_string("scalar(2.0)")

# Create input
input_stream = Stream(
    data=jnp.array([[1.0, 2.0, 3.0]]),
    dim_labels=(),
    chunk_size=1
)

# Compute Jacobian
jacobian = compute_jacobian(circuit, [input_stream])

print(jacobian)
# [[2. 0. 0.]
#  [0. 2. 0.]
#  [0. 0. 2.]]

The diagonal matrix shows each output depends only on its corresponding input, scaled by 2.

API Reference

def compute_jacobian(
    circuit: Circuit,
    inputs: List[Stream],
    output_idx: int = 0,
    input_idx: int = 0,
) -> jnp.ndarray

Parameters:

• circuit: Circuit to analyze
• inputs: Input streams
• output_idx: Which output stream (for multi-output circuits)
• input_idx: Which input stream (for multi-input circuits)

Returns:

• Jacobian matrix of shape (output_size, input_size)

Examples

Example 1: Addition Circuit

circuit = Circuit.from_string("add")

input1 = Stream(data=jnp.array([[1.0, 2.0]]), dim_labels=(), chunk_size=1)
input2 = Stream(data=jnp.array([[3.0, 4.0]]), dim_labels=(), chunk_size=1)

# Jacobian w.r.t. first input
jacobian1 = compute_jacobian(circuit, [input1, input2], input_idx=0)
# [[1. 0.]
#  [0. 1.]]
# Each output element depends equally on corresponding input1 element

# Jacobian w.r.t. second input
jacobian2 = compute_jacobian(circuit, [input1, input2], input_idx=1)
# [[1. 0.]
#  [0. 1.]]
# Same pattern for input2

Example 2: Composed Circuit

# Circuit: 2x + 3
circuit = Circuit.from_string(
    "composition(monoidal(scalar(2.0), const(3.0)), add)"
)

input_stream = Stream(data=jnp.array([[1.0, 2.0, 3.0]]), dim_labels=(), chunk_size=1)

jacobian = compute_jacobian(circuit, [input_stream])
# [[2. 0. 0.]
#  [0. 2. 0.]
#  [0. 0. 2.]]
# The constant term doesn't affect the Jacobian (derivative of constant is 0)

Example 3: Integration Circuit

circuit = Circuit.from_string("register(x)")

input_stream = Stream(
    data=jnp.array([[1.0, 2.0, 3.0, 4.0]]),
    dim_labels=("x",),
    chunk_size=1
)

jacobian = compute_jacobian(circuit, [input_stream])
# Shows how integration couples adjacent coefficients

Sensitivity Analysis

Identifying Important Inputs

The Jacobian norm for each input dimension indicates its importance:

# Compute Jacobian
jacobian = compute_jacobian(circuit, [input_stream])

# Importance of each input dimension
input_importance = jnp.linalg.norm(jacobian, axis=0)

print("Input importance:", input_importance)
# Higher values = more important inputs

Stability Analysis

The condition number indicates sensitivity to perturbations:

# Condition number (ratio of largest to smallest singular value)
cond = jnp.linalg.cond(jacobian)

if cond > 100:
    print("Warning: Circuit is sensitive to input perturbations")

Visualization

Jacobian Heatmap

import matplotlib.pyplot as plt

jacobian = compute_jacobian(circuit, [input_stream])

plt.figure(figsize=(8, 6))
plt.imshow(jacobian, cmap='RdBu', aspect='auto')
plt.colorbar(label='∂output/∂input')
plt.xlabel('Input dimension')
plt.ylabel('Output dimension')
plt.title('Jacobian Matrix')
plt.show()

Sensitivity Plot

# Compute sensitivity at different input values
sensitivities = []
x_values = jnp.linspace(-2, 2, 50)

for x in x_values:
    input_stream = Stream(
        data=jnp.array([[float(x)]]),
        dim_labels=(),
        chunk_size=1
    )
    jacobian = compute_jacobian(circuit, [input_stream])
    sensitivities.append(float(jacobian[0, 0]))

plt.plot(x_values, sensitivities)
plt.xlabel('Input value')
plt.ylabel('Sensitivity (∂output/∂input)')
plt.title('Sensitivity vs. Input')
plt.show()

Multi-Output Circuits

For circuits with multiple outputs, compute Jacobians for each:

# Parallel circuit with two outputs
circuit = Circuit.from_string("monoidal(scalar(2.0), scalar(3.0))")

input1 = Stream(data=jnp.array([[1.0]]), dim_labels=(), chunk_size=1)
input2 = Stream(data=jnp.array([[1.0]]), dim_labels=(), chunk_size=1)

# Jacobian for first output w.r.t. first input
j00 = compute_jacobian(circuit, [input1, input2], output_idx=0, input_idx=0)
# [[2.]]

# Jacobian for second output w.r.t. second input
j11 = compute_jacobian(circuit, [input1, input2], output_idx=1, input_idx=1)
# [[3.]]

Use Cases

1. Parameter Tuning

Use Jacobians to understand how parameters affect outputs:

# How does changing scalar affect output?
for scale in [1.0, 2.0, 3.0]:
    circuit = Circuit.from_string(f"scalar({scale})")
    jacobian = compute_jacobian(circuit, [input_stream])
    print(f"scale={scale}: sensitivity={jacobian[0,0]:.2f}")

2. Debugging Circuits

Verify circuits behave as expected:

# Fundamental theorem: composition(register, deregister) ≈ identity
circuit = Circuit.from_string("composition(register(x), deregister(x))")
jacobian = compute_jacobian(circuit, [input_stream])

# Should be approximately identity matrix (up to truncation)
identity_error = jnp.linalg.norm(jacobian - jnp.eye(len(jacobian)))
print(f"Deviation from identity: {identity_error:.6f}")

3. Optimization Guidance

Use Jacobians to guide parameter optimization:

# Compute gradient of loss w.r.t. parameters via chain rule
# dL/dθ = dL/dy * dy/dθ

output_jacobian = compute_jacobian(circuit, [input_stream])
loss_gradient = 2 * (predicted - target)  # MSE gradient

param_gradient = loss_gradient @ output_jacobian

Next Steps

• Gradient Optimization - Parameter optimization
• Circuits - Circuit basics
• API Reference - Complete API documentation