splax

Differentiable 3D gaussian splatting for JAX, with rasterization kernels written in NVIDIA Warp.

splax renders and trains 3D gaussian splats inside JAX. Projection, rasterization, and their backward passes run as Warp kernels wired into JAX through FFI custom calls under jax.custom_vjp, so rendering composes with jax.vmap, jax.grad, and jax.jit. No system CUDA toolchain is required.

Hero example

import jax.numpy as jnp
import splax

means, scales, quats, colors, opacities = splax.io.load_ply("scene.ply")
img = splax.inference.render(
    means, scales, quats, colors, opacities, viewmat=viewmat,
    background=jnp.ones(3), img_shape=(H, W), f=(fx, fy),
)  # (H, W, 3)

Two render entry points

splax exposes two renderers that share every Warp kernel and differ only in their JAX-level wrapping.

• splax.inference.render is the pure, grad-free forward path. Use it to serve a baked scene.
• splax.training.render, aliased as splax.render, is the differentiable path. Use it with jax.grad to fit gaussians.

Where to go next

• Installation covers the pip install, GPU requirements, and the pixi developer setup.
• Quickstart walks through rendering a scene, batching with jax.vmap, and taking a gradient.
• User Guide documents rendering, training, batching, and PLY IO.
• API Reference is generated from the source.