Gaussian splatting in PyTorch

In gfx-gaussian-splatting you rebuilt the 3DGS rasterizer the way the CUDA code does it: one thread per Gaussian, a radix sort, one thread per pixel. Research code looks different. Before anyone writes a kernel, the method is a few dozen lines of batched PyTorch: every Gaussian is a row of a tensor, every stage is a broadcast or a batched matmul, and autograd supplies the backward pass. gsplat keeps exactly such a reference next to its kernels (gsplat/cuda/_torch_impl.py, _math.py) and tests the CUDA against it; most new splatting ideas are prototyped in that style first.

This lesson builds that reference, gsplat's conventions first, and points out where the original implementation (graphdeco-inria/gaussian-splatting with its diff-gaussian-rasterization kernel, "3DGS" below) does things differently. You need autograd and parameter activations (py-torch-autograd) and camera matrices with OpenCV axes (py-camera-rays).

1. What the optimiser stores: raw parameters and activations

A splat has a mean, a rotation, three scales, an opacity and a colour. The optimiser never sees the physical values. It updates unconstrained tensors and the renderer maps them to legal values on every forward pass:

The activations live in the training code, not in the rasterizer: gsplat's rasterization() takes torch.exp(splats["scales"]) and torch.sigmoid(splats["opacities"]); 3DGS wraps them in get_scaling, get_opacity and get_rotation. The .ply files both write store the raw values — the reason a viewer must apply exp and sigmoid on load.

import torch

# Raw, unconstrained parameters: what the optimiser actually stores and updates.
log_scales = torch.tensor([[-3.0, -3.0, -5.0]])     # metres, in log space
opacity_logits = torch.tensor([0.0, 4.0, -4.0])
quats = torch.tensor([[2.0, 0.0, 0.0, 0.0]])        # (w, x, y, z), not unit

print("scales   ", torch.exp(log_scales).numpy().round(4))
print("opacities", torch.sigmoid(opacity_logits).numpy().round(4))
print("unit quat", torch.nn.functional.normalize(quats, dim=-1).numpy())

# Why log space: an Adam step changes the raw value by ~lr, i.e. the scale by a
# *factor* exp(±lr) — the same relative change for a 1 mm splat and a 1 m one.
for s in (0.001, 1.0):
    raw = torch.tensor(s).log()
    print(f"scale {s:5}: one step of 0.01 -> {torch.exp(raw + 0.01).item():.6f}")

Adam's step size is roughly lr per parameter regardless of the gradient's magnitude, so in log space every splat — a millimetre of hair or a metre of sky — grows or shrinks by the same 1 %. With linear scales the small splats would flip negative and the big ones would barely move. The logit plays the same role for opacity, and lets 3DGS's periodic "opacity reset" cap every splat at 0.01 in one line (inverse_sigmoid(torch.min(opacity, 0.01))).

2. Σ = R S Sᵀ Rᵀ, by broadcasting

The covariance must stay symmetric positive semi-definite while the optimiser pushes on it, so 3DGS builds it from a rotation and scales: M = R S, Σ = M Mᵀ. In PyTorch nobody forms S as a matrix. Multiplying R by scales[..., None, :] — shape [..., 1, 3] — scales column j by s_j, which is R @ diag(s):

import math
import torch

def quat_to_rotmat(quats):                          # [..., 4] (w, x, y, z) -> [..., 3, 3]
    q = quats / quats.norm(dim=-1, keepdim=True).clamp_min(1e-12)
    w, x, y, z = torch.unbind(q, dim=-1)
    R = torch.stack([
        1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y),
        2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x),
        2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)], dim=-1)
    return R.reshape(quats.shape[:-1] + (3, 3))

h = math.radians(45) / 2
quats = torch.tensor([[math.cos(h), 0.0, 0.0, math.sin(h)]])   # 45° about z
scales = torch.tensor([[2.0, 0.5, 1.0]])

R = quat_to_rotmat(quats)                           # [1, 3, 3]
M = R * scales[..., None, :]                        # [1,3,3] * [1,1,3]: column j times s_j
print(torch.allclose(M, R @ torch.diag_embed(scales)))
cov = M @ M.transpose(-1, -2)                       # R S Sᵀ Rᵀ
print(cov[0].numpy().round(4))

wrong = R * scales[..., :, None]                    # [1,3,1]: row i times s_i = S R
print((wrong @ wrong.transpose(-1, -2))[0].numpy().round(4))

batch = torch.randn(2, 5, 4)                        # any leading dims
print(quat_to_rotmat(batch).shape, quat_to_rotmat(torch.zeros(4)).diagonal())

The 2 × 0.5 ellipse turned 45° gets equal variances (4 + 0.25)/2 = 2.125 and a covariance (4 − 0.25)/2 = 1.875 — the same numbers as the C++ lesson. The wrong broadcast gives diag(4, 0.25, 1): right eigenvalues, no rotation, a bug that a test on the spectrum alone never catches.

Three habits from gsplat's _math.py worth copying:

• unbind + stack + reshape(quats.shape[:-1] + (3, 3)) writes the nine entries once and works for any batch shape — [N], [C, N], [B, C, N].
• Normalise with a clamp (F.normalize uses max(‖q‖, 1e-12)): the zero quaternion becomes the identity instead of NaN — the last line above. Gradients push some quaternions towards zero during training.
• w first. 3DGS's rot_0..3, gsplat's quats and nerfstudio are (w, x, y, z). SciPy's Rotation.from_quat, Eigen's coefficient storage and glTF are (x, y, z, w). Convert explicitly at every boundary.

(3DGS itself writes L = R @ S with an explicit S = diag(s) in build_scaling_rotation; its CUDA computeCov3D spells it glm::mat3 M = S * R and is still right because glm is column-major — see gfx-gaussian-splatting.)

3. Into the camera and onto the screen

gsplat takes a world-to-camera viewmat [4, 4] (OpenCV axes: +x right, +y down, +z forward) and intrinsics K [3, 3]. A linear map x ↦ R x takes a covariance to R Σ Rᵀ, so _world_to_cam is two lines: means @ R.T + t and R @ covars @ R.T. (3DGS instead takes the FoV and a full projection matrix built with getProjectionMatrix, assumes the principal point at the centre, and transposes everything because glm is column-major.)

Perspective isn't affine, so EWA splatting (Zwicker et al. 2002) linearises it at the mean t = (x, y, z): the projection (fx·x/z + cx, fy·y/z + cy) has the 2 × 3 Jacobian

J = [[fx/z, 0, −fx·x/z²], [0, fy/z, −fy·y/z²]],

and Σ′ = J W Σ Wᵀ Jᵀ (W = the view rotation, already applied above). gsplat's _persp_proj first clamps x/z and y/z to the view frustum widened by 30 % of the half-FoV on each side — the linearisation is terrible far off-axis, and a splat just outside the view would otherwise smear across the screen — but only for J; the 2D mean stays where the splat really is.

import torch

W, H = 64, 48
K = torch.tensor([[100.0, 0.0, 32.0], [0.0, 100.0, 24.0], [0.0, 0.0, 1.0]])
means_c = torch.tensor([[0.4, -0.2, 2.0], [3.0, 0.0, 2.0]])   # camera space, OpenCV axes
covars_c = 0.01 * torch.eye(3).expand(2, 3, 3)                 # round, σ = 0.1

tx, ty, tz = torch.unbind(means_c, dim=-1)                     # each [N]
fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
tan_fovx, tan_fovy = 0.5 * W / fx, 0.5 * H / fy
jx = tz * torch.clamp(tx / tz, -(cx / fx + 0.3 * tan_fovx), (W - cx) / fx + 0.3 * tan_fovx)
jy = tz * torch.clamp(ty / tz, -(cy / fy + 0.3 * tan_fovy), (H - cy) / fy + 0.3 * tan_fovy)
O = torch.zeros_like(tz)
J = torch.stack([fx / tz, O, -fx * jx / tz**2,
                 O, fy / tz, -fy * jy / tz**2], dim=-1).reshape(-1, 2, 3)   # [N, 2, 3]
cov2d = J @ covars_c @ J.transpose(-1, -2)                     # [N, 2, 2]
means2d = torch.stack([fx * tx / tz + cx, fy * ty / tz + cy], dim=-1)
print("clamped x/z ", [round(v, 3) for v in (jx / tz).tolist()], " true", (tx / tz).tolist())
print("means2d     ", means2d.tolist())
print("cov2d[0]    ", cov2d[0].tolist())

cov2d = cov2d + 0.3 * torch.eye(2)                             # eps2d low-pass
a, b, c = cov2d[:, 0, 0], cov2d[:, 0, 1], cov2d[:, 1, 1]
det = (a * c - b * b).clamp_min(1e-10)
conics = torch.stack([c / det, -b / det, a / det], dim=-1)
print("conic[0]    ", [f"{v:.5f}" for v in conics[0].tolist()])
print("gsplat radii", torch.ceil(3.33 * torch.sqrt(torch.stack([a, c], -1)))[0].int().tolist())
mid = 0.5 * (a + c)
lam = mid + torch.sqrt(torch.clamp(mid * mid - det, min=0.1))
print("3DGS radius ", int(torch.ceil(3 * torch.sqrt(lam))[0]))

The second Gaussian sits at x/z = 1.5, far outside the 0.32 half-width of this view: its J is built at x/z ≈ 0.416 (the frustum edge 0.32 plus 0.3 × 0.32), while its mean projects honestly to x = 182 px. Everything is [N]-batched: unbind the coordinates, compute each entry as an [N] tensor, stack and reshape into [N, 2, 3], and let @ broadcast the matrix products over N.

4. The low-pass filter, the conic and the radius

cov2d + 0.3 * I (gsplat's eps2d) convolves every splat with a screen-space Gaussian of variance 0.3 px²: sub-pixel splats grow to about a pixel, they stop aliasing, and Σ′ stays invertible. Everything downstream uses the dilated matrix. Its inverse, the conic (A, B, C) = (c, −b, a)/det, is all the rasterizer needs: the exponent at offset d is σ = ½(A dx² + C dy²) + B dx dy.

How far to look is where the codebases differ:

• 3DGS: one radius ⌈3·√λ_max⌉ from the eigenvalues of Σ′ (with the term under the root floored at 0.1) — 16 px above.
• gsplat: per-axis radii ⌈3.33·√Σ′ₓₓ⌉, ⌈3.33·√Σ′ᵧᵧ⌉ — the exact extent of the 3.33σ ellipse along x and y, tighter for thin axis-aligned splats — [18, 17] above. When opacities are passed to the projection, gsplat also shrinks the extent to where the splat fades below 1/255: min(3.33, √(2 ln(255·o))) standard deviations.
• With rasterize_mode="antialiased" gsplat also multiplies opacity by the compensation √(det Σ′ / det(Σ′ + 0.3 I)) — Mip-Splatting's 2D filter, which keeps a tiny splat's total energy constant instead of fattening it.

In a dense renderer like this skill's, the radius only culls (radii 0 = not drawn) — every remaining splat is evaluated at every pixel and the 1/255 threshold does the rest. In a tiled renderer it decides which tiles a splat is duplicated into, so a tighter bound is directly fewer sort keys and less blending work.

5. Compositing as tensor ops

The CUDA kernel walks each pixel's sorted list in a loop:

for (each splat in depth order) {
    float sigma = 0.5f * (A * dx * dx + C * dy * dy) + B * dx * dy;
    float alpha = min(0.99f, opac * exp(-sigma));
    if (sigma < 0.f || alpha < 1.f / 255.f) continue;
    float next_T = T * (1.f - alpha);
    if (next_T <= 1e-4f) { done = true; break; }    // gsplat; 3DGS: test_T < 0.0001f
    pix += color * alpha * T;
    T = next_T;
}

PyTorch has no per-pixel loops worth writing. Instead: sort the M visible splats once (torch.argsort(depths, stable=True) — one global sort of M depths replaces the per-tile ones), evaluate them all at every pixel centre as one [M, H, W] tensor, and turn the loop's running product into a cumulative product along the M axis.

import torch

# Dense rasterization: every splat at every pixel centre, as one [M, H, W] tensor.
H, W = 4, 6
means2d = torch.tensor([[2.0, 1.5], [3.5, 2.0]])       # already sorted nearest first
conics = torch.tensor([[0.5, 0.0, 0.5], [0.3, 0.1, 0.4]])
opac = torch.tensor([0.9, 0.7])

xs = torch.arange(W) + 0.5                             # gsplat: pixel centres at +0.5
ys = torch.arange(H) + 0.5
dx = means2d[:, 0, None, None] - xs                    # [M, 1, W]
dy = means2d[:, 1, None, None] - ys[:, None]           # [M, H, 1]
A, B, C = (conics[:, k, None, None] for k in range(3))
sigma = 0.5 * (A * dx**2 + C * dy**2) + B * dx * dy    # broadcasts to [M, H, W]
alpha = torch.clamp_max(opac[:, None, None] * torch.exp(-sigma), 0.99)
alpha = torch.where((sigma >= 0) & (alpha >= 1 / 255), alpha, 0.0)
print(dx.shape, dy.shape, sigma.shape)
print((alpha[0] * 100).round().int())

Note dx is [M, 1, W] and dy is [M, H, 1]: broadcasting builds the full grid only when they meet. The first splat sits on the edge between pixels (1, 1) and (1, 2) — centres 0.5 px either side — so those two get the same, largest alpha.

Pixel centres. gsplat samples pixel (i, j) at (j + 0.5, i + 0.5), so the image spans [0, W) and cx = W/2 is its centre. The 3DGS kernel samples at (j, i) and its ndc2Pix puts the principal point at (W − 1)/2 — consistent on its own, half a pixel off when you mix the two.

Then the transmittance, along M. torch.cumprod is inclusive (it includes element i at position i); the splat at i must be weighted by the transmittance in front of it, so shift by one:

import torch

alpha = torch.tensor([0.6, 0.5, 0.9, 0.8])            # one pixel, nearest splat first
inclusive = torch.cumprod(1 - alpha, dim=0)
T = torch.cat([torch.ones(1), inclusive[:-1]])         # exclusive: T_0 = 1
print("inclusive", [round(v, 4) for v in inclusive.tolist()])
print("T        ", [round(v, 4) for v in T.tolist()])
w = alpha * T
print("weights  ", [round(v, 4) for v in w.tolist()], " sum", round(w.sum().item(), 4))

# Early termination, vectorised: keep a splat only if T after it stays above 1e-4.
alpha = torch.tensor([0.98, 0.98, 0.98, 0.98])
T = torch.cat([torch.ones(1), torch.cumprod(1 - alpha, 0)[:-1]])
keep = T * (1 - alpha) > 1e-4
print("T        ", [f"{v:.1e}" for v in T.tolist()], "keep", keep.tolist())
print("acc      ", f"{(alpha * T * keep).sum().item():.5f}", "= 1 - T_final =", f"{1 - T[2].item():.5f}")

The loop's break becomes a mask: T only ever shrinks, so "T after this splat

1e-4" is true for a prefix of the list and false from the terminating splat on — exactly the splats the kernel adds. The third splat would take T from 4·10⁻⁴ to 8·10⁻⁶, so like the kernel we stop without it; the pixel's alpha is 1 − 4·10⁻⁴. (NeRF code does the same shift: tf.math.cumprod(..., exclusive=True) in bmild/nerf, torch.cumprod(torch.cat([ones, 1. - alpha + 1e-10], -1), -1)[:, :-1] in nerf-pytorch — see py-nerf-rendering.)

Finally image = einsum("mhw,mc->hwc", weights, colors) + (1 − Σw)·background: C channels of any kind (RGB, depth, features) cost the same einsum.

6. Differentiable end to end — and NaN-free

Everything above is a tensor op, so loss.backward() reaches every raw parameter: through exp and sigmoid, the quaternion normalisation, both matmuls, J, the conic inverse, exp(−σ), cumprod and the einsum. The sort contributes no gradient (it is an integer permutation) but indexing with it does; the radii are integers and only select. Three numerical details decide whether training survives:

• The α cap. clamp_max(·, 0.99) keeps 1 − α ≥ 0.01. The CUDA backward recovers Tᵢ by dividing the final T by (1 − αⱼ) on its way back to front; PyTorch's cumprod backward also divides when it can. A splat with α = 1 would make that 0/0. (Older gsplat capped at 0.999; current gsplat and 3DGS use 0.99.)
• The det clamp. det.clamp_min(1e-10) and the 0.3 dilation keep the conic finite even for a needle-thin splat seen end-on.
• Masked inputs, not masked outputs. A Gaussian behind the camera or at z = 0 is culled — but if you still divide by its z and only zero the result, the forward pass is clean and the backward pass is not:

import torch

z = torch.tensor([2.0, 0.0, -1.0], requires_grad=True)   # depths; only the first is valid
valid = z > 0.01

# Mask the output: the forward pass is clean...
u = torch.where(valid, 100.0 / z, 0.0)
u.sum().backward()
print("masked output ", u.detach().tolist(), "grad", z.grad.tolist())

# ...sanitise the input instead (the "double where"): nothing non-finite is ever computed.
z.grad = None
z_safe = torch.where(valid, z, 1.0)
u = torch.where(valid, 100.0 / z_safe, 0.0)
u.sum().backward()
print("masked input  ", u.detach().tolist(), "grad", z.grad.tolist())

torch.where routes a zero gradient into the masked-out branch, and the division's backward multiplies it by −100/z² = −inf: 0 · inf = NaN. One NaN in a parameter's gradient makes Adam's moment estimates NaN, and the splat is gone for good. gsplat's CUDA projection simply returns for culled Gaussians, so their gradients are exactly 0; a PyTorch port must replace their inputs before anything divides.

Test differentiable code in float64 with torch.autograd.gradcheck: it compares autograd against finite differences, and float32's 7 digits are too few for that. That means every constant you create — torch.eye(2), torch.arange(W), the stand-in depth — must follow the input's dtype (dtype=x.dtype, x.new_tensor).

7. Fitting splats with Adam

Once rendering is a differentiable function, "optimising a radiance field" is the same loop as training any model. A 2D toy with isotropic splats over a white background:

import torch

torch.manual_seed(0)
H, W, N = 16, 16, 6
ys, xs = torch.meshgrid(torch.arange(H) + 0.5, torch.arange(W) + 0.5, indexing="ij")

def render(means, log_sigma, logits, colors):
    d2 = (means[:, 0, None, None] - xs) ** 2 + (means[:, 1, None, None] - ys) ** 2
    var = torch.exp(2 * log_sigma)[:, None, None]
    alpha = (torch.sigmoid(logits)[:, None, None] * torch.exp(-0.5 * d2 / var)).clamp_max(0.99)
    T = torch.cat([torch.ones_like(alpha[:1]), torch.cumprod(1 - alpha, 0)[:-1]])
    w = alpha * T
    return torch.einsum("nhw,nc->hwc", w, colors) + (1 - w.sum(0))[..., None]   # white background

target = torch.ones(H, W, 3)
target[4:12, 4:12] = torch.tensor([0.9, 0.2, 0.1])                   # a red square

params = {"means": torch.rand(N, 2) * W, "log_sigma": torch.full((N,), 1.0),
          "logits": torch.zeros(N), "colors": torch.full((N, 3), 0.5)}
for p in params.values():
    p.requires_grad_(True)
opt = torch.optim.Adam(params.values(), lr=0.1)
for step in range(201):
    loss = (render(**params) - target).abs().mean()
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step % 50 == 0:
        print(f"step {step:3d}  L1 {loss.item():.4f}")

The loss falls by about 3× and then stalls: six round blobs can't make a sharp square. Real 3DGS training fixes that with adaptive density control — clone small splats with large view-space positional gradients, split large ones, prune transparent ones, periodically reset opacities — and uses per-parameter learning rates (means decayed from 1.6·10⁻⁴ times the scene extent, scales 5·10⁻³, opacities 2.5·10⁻² in the latest code, quaternions 10⁻³) and the loss (1 − λ)·L1 + λ·(1 − SSIM) with λ = 0.2 (py-image-metrics). gsplat packages the densification heuristics as DefaultStrategy and MCMCStrategy.

8. Dense vs tiled, and what to say in an interview

The dense rasterizer is O(M·H·W) in time and memory — autograd keeps the [M, H, W] alphas for the backward pass. That's fine for a 64 × 48 test image with a few hundred splats, hopeless for a million splats at 1080p (≈ 2·10¹² entries). The production renderers bin splats into 16 × 16 tiles so each pixel only walks its tile's list (gfx-gaussian-splatting), stop each pixel at T ≤ 10⁻⁴, and write the backward pass by hand, back to front. The PyTorch version is still the one you write first — to try an idea, and to test the kernel against.

In an interview — "implement the forward pass of 3DGS in PyTorch" is a realistic whiteboard question for this job. Be ready to say: raw parameters and their activations (log-scale, logit-opacity, normalised w-first quaternion); Σ = R S Sᵀ Rᵀ as R * s[..., None, :]; R Σ Rᵀ into the camera; the Jacobian J, its clamp, and Σ′ = J W Σ Wᵀ Jᵀ; the 0.3 dilation before the inverse; radii; depth sort; alpha with the 1/255 and 0.99 thresholds; exclusive cumprod; early termination as a mask; and why you sanitise culled inputs before dividing. Then the follow-ups: how would you make it scale? (tiles, a fused kernel, packed sparse intersections); where do gsplat and the original differ? (pixel centres, viewmat + K vs FoV, radius rule, ≤ vs <, the antialiased mode).