splax - GitLite

KO | EN

gitlite — search

#cpp #pdf #bash #vue #github #awesome #typescript #nextjs #electron #nlp

splax

(★ 9)

Gaussian Splatting in JAX

File Explorer

ZIP Download

- - docs.yml
  - ruff.yml
  - ty.yml
- bench_backward.py
- bench_forward.py
- suite_report.py
- - index.md
- - gen_logo.py
  - logo.png
  - logo.svg
  - logo_emerges_dark.gif
  - logo_emerges_light.gif
  - logo_tiles.json
- - install.md
  - quickstart.md
- - mathjax.js
- - batching.md
  - colmap.md
  - io.md
  - rendering.md
  - training.md
  - viewer.md
- examples.md
- gen_ref_pages.py
- index.md
- compose_splats.py
- distortion.py
- render_scene.py
- viewer_demo.py
- Apache-2.0.txt
- __init__.py
- colmap.py
- logo_editor.py
- logo_splat.py
- train_colmap.py
- train_lego.py
- - __init__.py
  - _kernels.py
  - _project.py
- - __init__.py
  - _kernels.py
  - _rasterize.py
- __init__.py
- _batching.py
- _cache.py
- _harmonics.py
- _intersect.py
- _render.py
- io.py
- mcmc.py
- py.typed
- utils.py
- viewer.py
- - test_batched_training.py
  - test_colmap_loader.py
  - test_examples.py
  - test_harmonics.py
  - test_io.py
  - test_lego_regression.py
  - test_mcmc.py
  - test_project.py
  - test_rasterize.py
  - test_render.py
  - test_render_grad.py
  - test_transforms.py
- - test_cache.py
  - test_colmap_init.py
  - test_exposure.py
  - test_harmonics.py
  - test_io.py
  - test_mcmc.py
  - test_project.py
  - test_rasterize.py
  - test_transforms.py
  - test_utils.py
  - test_viewer.py
- _gsplat.py
- conftest.py
- utils.py
.gitignore
LICENSE
pixi.lock
properdocs.yml
pyproject.toml
README.md

# Use via CDN

jsDelivr

jsDelivr serves any public GitHub repository as a CDN with zero setup. Pick a version and a file to get a ready-to-paste link and snippet.

Version

File

Link

Example

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

<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">

**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)](https://raw.githubusercontent.com/learnsyslab/splax/main/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.

## Installation

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

```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 (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](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)).