---
id: sax
version: "0.18.2"
license: Apache Software License
license_treatment: permissive
maintenance: active
---

# sax — Autograd and XLA for S-parameters

License: permissive · Maintenance: active · Downloads: 136.3K/mo

## What it is and what it does

SAX wraps JAX to provide a functional, dictionary-based framework for simulating and optimizing circuits described by S-parameters (scattering matrices). It was built for photonic integrated circuits but works with any frequency-domain circuit model. Rather than defining custom data structures, SAX stays close to JAX's functional paradigm: you write component models as functions returning S-dictionaries, compose them into circuits via netlists, and optimize them using JAX's autodiff and XLA compilation.

The package handles the boilerplate of circuit composition and parameter passing, allowing you to define a directional coupler, waveguide, or other component once and then combine them into larger systems like Mach-Zehnder interferometers. You can then sweep parameters (wavelength, coupling strength, length) across ranges and plot transmission spectra or use JAX's optimization tools to tune component parameters for a target response.

Use it for:

- Design and simulate photonic integrated circuits by composing waveguides, couplers, and other optical components into larger systems.
- Optimize photonic circuit parameters (coupling ratios, waveguide lengths, phase shifts) to achieve target transmission or reflection spectra.
- Perform frequency-domain S-parameter analysis on arbitrary circuits without writing custom matrix algebra.
- Leverage JAX autodiff to compute gradients for circuit optimization or sensitivity analysis across wavelength ranges.
- Combine circuit simulation with machine learning workflows via JAX's ecosystem (Optax, Flax, etc.).

## Worth the install?

AI-flagged interpretation of the facts above — verify before relying.

SAX is a JAX-based simulator and optimizer for S-parameter circuits, primarily designed for photonic integrated circuits but applicable to any frequency-domain circuit simulation using scattering parameters.

Yes, if you work with S-parameter circuit simulation or photonic design. SAX is actively maintained, has no known vulnerabilities, and offers a clean functional interface built on JAX's autodiff and XLA compilation. The dependency footprint is large but manageable. Install with caution if your environment has strict dependency constraints; otherwise, it is a solid choice for frequency-domain circuit work.

## Install

pip install sax
uv add sax
poetry add sax

## Installing sax

Before you install: Low install friction with a pure-wheel distribution. Active maintenance (66 days since last release) and support for current Python versions (3.11–3.14). Requires 19 runtime dependencies including JAX, NumPy, and scientific libraries; installation is straightforward but the dependency stack is substantial.

License in practice: Apache-2.0 permissive license allows commercial and private use with minimal restrictions; you must include a copy of the license and state significant changes, but there are no copyleft obligations.

Quickstart:

pip install sax

import sax
import jax.numpy as jnp

def coupler(coupling=0.5):
    kappa = coupling**0.5
    tau = (1-coupling)**0.5
    return sax.reciprocal({
        ("in0", "out0"): tau,
        ("in0", "out1"): 1j*kappa,
        ("in1", "out0"): 1j*kappa,
        ("in1", "out1"): tau,
    })

result = coupler(coupling=0.3)

Requires Python >=3.11.0; JAX installation may require additional system dependencies depending on your platform (CPU vs. GPU/TPU support).

Verify before relying:

- Whether the package's 19 dependencies can be installed together without version conflicts in typical environments.
- Performance characteristics and scalability limits for large circuit topologies or high-dimensional parameter sweeps.

## Package facts

- License: Apache Software License (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 136.3K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

s-parameter circuit simulator, photonic circuit simulation, jax-based optimization, frequency domain circuit analysis, autograd circuit modeling, photonic integrated circuit design, s-matrix simulation, photonics, circuit-simulation, jax-based

[View on SkillFed](https://skillfed.io/packages/sax) · [View on PyPI](https://pypi.org/project/sax/)