jaxstream

High-frequency signal streaming from JAX models via WebSocket.

Status: Published ahead of end-to-end validation. Treat 0.x as provisional — the API is intended to work as advertised, but hasn't been exercised too much yet.

Why

When running JAX models (especially on GPU), extracting intermediate values from JIT-compiled computation graphs is non-trivial. JAX's io_callback is the sanctioned way to perform side effects from inside a traced function — and livestream() is designed to be passed directly to it.

This gives you a drop-in way to tap into any point in a model's forward pass and stream arrays to external consumers over WebSocket in real time.

Installation

uv add jaxstream

Quick start

1. Register a stream with a transform

Transforms process raw content before it's sent to clients. The "default" stream (identity transform) is always available.

from jaxstream import register_stream

def normalize(raw):
    """Normalize each array to [0, 1]."""
    return {k: (v - v.min()) / (v.max() - v.min()) for k, v in raw.items()}

register_stream("normalized_actions", normalize)

2. Stream from inside a JAX model

from jax.experimental import io_callback
from jaxstream import livestream

# Inside a JIT-compiled function:
io_callback(
    livestream,
    (),  # no return value
    {"normalized_actions": {"positions": position_array, "velocities": vel_array}}
)

livestream() automatically starts a WebSocket server on port 8765 on first call. Each key in the dict is a stream name; its value is passed through the registered transform before serialization.

3. Connect a client

import asyncio
import msgpack
import websockets

async def listen():
    async with websockets.connect("ws://localhost:8765") as ws:
        while True:
            data = msgpack.unpackb(await ws.recv(), timestamp=3)
            print(data)

asyncio.run(listen())

Messages are serialized with msgpack. JAX and numpy arrays are automatically converted via .tolist(). Each message includes sourced_time (when livestream() was called) and processed_time (when the message was sent).

API

livestream(streams: dict[Livestream, Content])

Broadcast data to all connected WebSocket clients. Manages a global WebsocketStreamer singleton. Designed to be used with jax.experimental.io_callback.

register_stream(name, transform, validate_raw=None, validate=None)

Register a named stream with a transform function. Optionally attach validation functions for raw input and processed output.

process(raw, stream="default")

Apply a registered stream's transform to content. Used internally by the streamer, but available for testing transforms independently.

WebsocketStreamer(port=8765)

Low-level WebSocket server. Call .start() to run in a background thread, then .broadcast(stream_name, content) to queue messages.

Development

uv sync --extra dev
uv run pytest -v