---

splax

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

splax renders and trains 3D gaussian splats inside JAX. Batched rendering and training run through jax.vmap, jax.grad, and jax.jit, with no system CUDA toolchain needed.

Examples

Render a .ply scene from one camera.

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)

Batch over a stack of camera poses with jax.vmap. One batched kernel launch, not a Python loop.

import jax

frames = jax.vmap(lambda vm: splax.inference.render(
    means, scales, quats, colors, opacities,
    viewmat=vm, background=jnp.ones(3), img_shape=(H, W), f=(fx, fy),
))(viewmats)  # (B, H, W, 3)

Take gradients through the differentiable renderer with jax.grad. splax.render is splax.training.render and differentiates with respect to means, scales, quats, colors, opacities.

import jax

def loss(means, scales, quats, colors, opacities):
    img, _ = splax.render(
        means, scales, quats, colors, opacities,
        viewmat=viewmat, background=jnp.ones(3), img_shape=(H, W), f=(fx, fy),
    )
    return jnp.mean((img - target) ** 2)

grads = jax.grad(loss, argnums=(0, 1, 2, 3, 4))(means, scales, quats, colors, opacities)

To edit .ply scenes, we recommend superspl.at.

Documentation

Full documentation lives at learnsyslab.github.io/splax: installation, a quickstart, a user guide for rendering, training, batching, and IO, and the API reference.

Why

Gaussian splatting lives mostly in PyTorch and hand-written CUDA. splax puts it inside JAX so splat rendering composes with jax.vmap, jax.grad, and jax.jit and drops into research pipelines that already run on JAX, without leaving the ecosystem for the render step.

Architecture

The renderer is not pure JAX because the core of splatting does not map to XLA primitives. Rasterization is tile-binned with a data-dependent sort of gaussian-tile intersections, per-pixel early termination once transmittance saturates, and a memory-frugal backward that recomputes the blend instead of storing per-pixel state. splax implements the projection, rasterization, and their backward passes as Warp kernels and wires them into JAX through FFI custom calls under jax.custom_vjp. The kernels are batch-native: jax.vmap maps to a single batched launch (camera id folded into the sort key) rather than a sequential per-sample loop.

Relation to jaxsplat

splax started from jaxsplat as the reference and parity baseline. The rasterizer was rewritten from CUDA to Warp to drop the system toolchain, then extended with the feature and performance work below.

Improvements

Ported from gsplat and the papers behind it, which inspired most of the performance work (credit per item).

• Native multi-camera batched rendering, one launch with the camera id folded into the sort key (gsplat)
• Opacity-aware tight tile intersection (StopThePop, Speedy-Splat, gsplat #927)
• Packed 32-bit sort keys with quantized depth
• Persistent sort and bin scratch across frames (gsplat caching allocator design)
• Cooperative shared-memory tile blending with block-vote early exit (3DGS, gsplat)
• Fixed-budget MCMC training with static shapes, relocation plus covariance noise (gsplat MCMCStrategy, Kheradmand et al. 2024)
• Opacity and scale regularizers (gsplat mcmc preset)
• Progressive resolution fine-tuning (coarse-to-fine, 3DGS)
• Per-parameter Adam learning-rate schedules (gsplat, 3DGS)
• L1 plus D-SSIM photometric loss (3DGS, gsplat, via dm-pix)
• Camera pose gradients via jax.grad(loss, argnums=viewmat), dispatched by JAX symbolic zeros so a pose-only step skips the gaussian projection backward (gsplat projection backward)
• Batch-native backward passes, jax.vmap(jax.grad(render)) runs as one batched launch (gsplat)
• Batched training steps with sqrt-batch learning-rate scaling (gsplat batch_size and steps_scaler)
• Anti-aliased opacity compensation (Mip-Splatting, gsplat), depth regularization from COLMAP points (gsplat depth_loss), per-image exposure correction (gsplat appearance optimization), all opt-in

Installation

Requires an NVIDIA GPU and a CUDA-enabled JAX (jax[cuda12], pulled in as a dependency).

uv pip install "git+https://github.com/learnsyslab/splax"

Developer setup with pixi, which installs splax editable with the dev tooling:

git clone https://github.com/learnsyslab/splax.git
cd splax
pixi shell

Run the test suite with pixi run -e tests tests. The suite also runs against any splax installation that includes the tests extra. Checking a built distribution before a release therefore is:

pixi run -e dist build
uv venv /tmp/splax-check
source /tmp/splax-check/bin/activate
uv pip install "dist/splax-0.1.0-py3-none-any.whl[tests]"
cd tests && pytest .

The gsplat reference tests JIT-compile a CUDA extension on first use, which needs a CUDA 12.8 compiler toolchain that pip cannot provide. Run the commands inside pixi shell -e tests to use the toolchain from the tests environment, or provide a system CUDA 12.8 install.

License

MIT (see LICENSE). gsplat-derived portions are under Apache-2.0 (licenses/Apache-2.0.txt).