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.

Installation

pip install splax

splax needs an NVIDIA GPU and requires CUDA-enabled JAX. No system CUDA toolchain is required.

Examples

Render a .ply scene from one camera.

import jax.numpy as jnp
import splax

splats = splax.io.load_ply("scene.ply")  # means, log_scales, quats, sh_colors, logit_opacities
img, _ = splax.render(
    *splats, 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, rendered in one go rather than in a Python loop.

import jax

frames = jax.vmap(
    lambda vm: splax.render(
        *splats, viewmat=vm, background=jnp.ones(3), img_shape=(H, W), f=(fx, fy)
    )[0]
)(viewmats)  # (B, H, W, 3)

Take gradients through the renderer with jax.grad. splax.render differentiates with respect to means, log scales, quats, SH colors, logit opacities, the camera pose, and per-object rigid transforms.

import jax


def loss(means, log_scales, quats, sh_colors, logit_opacities):
    img, _ = splax.render(
        means,
        log_scales,
        quats,
        sh_colors,
        logit_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))(*splats)

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

Splatting does not map onto XLA primitives, so projection, rasterization, and their backward passes are Warp kernels called from JAX. Rendering therefore composes with jax.vmap, jax.grad, and jax.jit, and a batch renders in one go rather than looping in Python.

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 (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 (gsplat projection backward)
• Brown-Conrady lens distortion applied inside the projection
• Batch-native backward passes under jax.vmap(jax.grad(render)) (gsplat)
• Batched training steps with sqrt-batch learning-rate scaling (gsplat batch_size and steps_scaler)
• View-dependent color from spherical harmonics up to degree 3 (3DGS, gsplat)
• 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

Development

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).