Metadata-Version: 2.4
Name: splax
Version: 0.1.0
Summary: High-performance vectorized gaussian splatting rendering in JAX
Author: Martin Schuck
License-Expression: MIT
Project-URL: Homepage, https://github.com/learnsyslab/splax
Project-URL: Repository, https://github.com/learnsyslab/splax
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Education
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: licenses/Apache-2.0.txt
Requires-Dist: jax[cuda]>=0.10
Requires-Dist: imageio
Requires-Dist: plyfile
Requires-Dist: warp-lang>=1.14.0
Provides-Extra: viewer
Requires-Dist: viser>=1.0; extra == "viewer"
Provides-Extra: tests
Requires-Dist: splax[viewer]; extra == "tests"
Requires-Dist: pytest; extra == "tests"
Requires-Dist: pytest-markdown-docs; extra == "tests"
Requires-Dist: torch<2.11,>=2.8; extra == "tests"
Requires-Dist: gsplat>=1.5; extra == "tests"
Requires-Dist: optax>=0.2.8; extra == "tests"
Requires-Dist: dm-pix>=0.4.5; extra == "tests"
Requires-Dist: pycolmap>=4.1; extra == "tests"
Requires-Dist: matplotlib; extra == "tests"
Dynamic: license-file

<div align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/learnsyslab/splax/main/docs/assets/logo_emerges_dark.gif">
    <img src="https://raw.githubusercontent.com/learnsyslab/splax/main/docs/assets/logo_emerges_light.gif" alt="splax" width="560">
  </picture>
</div>

--------------------------------------------------------------------------------

<div align="center">

# splax

**Differentiable 3D gaussian splatting for JAX, with rasterization kernels written in [NVIDIA Warp](https://github.com/NVIDIA/warp).**

[![Python](https://img.shields.io/badge/python-3.11%2B-blue.svg)](https://www.python.org)
[![License](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE)
[![Ruff](https://github.com/learnsyslab/splax/actions/workflows/ruff.yml/badge.svg)](https://github.com/learnsyslab/splax/actions/workflows/ruff.yml)
[![Ty](https://github.com/learnsyslab/splax/actions/workflows/ty.yml/badge.svg)](https://github.com/learnsyslab/splax/actions/workflows/ty.yml)
[![Docs](https://github.com/learnsyslab/splax/actions/workflows/docs.yml/badge.svg)](https://learnsyslab.github.io/splax)

</div>

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.

```python
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.

```python
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.

```python
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](https://superspl.at/editor).

## Documentation

Full documentation lives at [learnsyslab.github.io/splax](https://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](https://github.com/yklcs/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](https://github.com/nerfstudio-project/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, where a pose-only step costs only the camera gradient (gsplat projection backward)
- 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`)
- 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[cuda]`, pulled in as a dependency).

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

Developer setup with [pixi](https://pixi.sh/), which installs splax editable with the dev tooling:

```sh
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:

```sh
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](LICENSE)). gsplat-derived portions are under Apache-2.0 ([licenses/Apache-2.0.txt](licenses/Apache-2.0.txt)).
