Skip to content

splax

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 are Warp kernels called from JAX, so rendering composes with jax.vmap, jax.grad, and jax.jit. No system CUDA toolchain is required.

Rendering a scene

import jax.numpy as jnp
import splax

SCENE = "https://huggingface.co/datasets/amacati/splax-test-data/resolve/main/scenes/lego.ply"
splats = splax.io.load_ply(splax.io.fetch(SCENE))
H, W, fx, fy = 400, 400, 400.0, 400.0
viewmat = splax.utils.look_at(jnp.array((0.0, -3.0, 1.0)), jnp.zeros(3), up=(0.0, 0.0, 1.0))
img, _ = splax.render(
    *splats, viewmat=viewmat, background=jnp.ones(3), img_shape=(H, W), f=(fx, fy)
)  # (H, W, 3)

Render entry point

splax.render handles the rendering and is differentiable with respect to the gaussian parameters, the camera pose, and per-object rigid transforms.

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.