Differentiate the EWA Jacobian through the FOV clamp

[manuel-rodriguezvillegas]

Problem

_project_kernel evaluates the EWA Jacobian at the camera-space position clamped to 1.3 times the half field of view, but _project_bwd_kernel rebuilt the Jacobian from the unclamped position ("the standard gsplat approximation"), and _proj_vjp differentiated that unclamped Jacobian. For a gaussian whose center lies beyond the clamp while its footprint still reaches into the image, the backward pass was therefore not the derivative of the forward pass. These are typically large gaussians close to the camera, such as a floor or a wall, and they can cover a large part of the image.

The existing gradient tests use scenes near the image center, where nothing is clamped, so they did not catch it.

Fix

• The clamp lives in two helpers, _fov_limits and _clamp_to_fov, which both kernels use, so the forward and backward passes cannot drift apart. The forward output is unchanged: it is bit-identical to 0.2.2 on 20,000 random gaussians placed up to 3 times beyond the limits.
• The backward kernel builds the Jacobian at the clamped position. _proj_vjp differentiates it as J = G(u, v) / z with u = clamp(x / z) and v = clamp(y / z), so the gradient along a clamped coordinate vanishes, as it does in the image. The center path is unchanged.
• The backward kernel and its FFI take the image size, from which they derive the same limits as the forward kernel.

Tests

• The new test_project_grad_clamped checks large gaussians beyond the limits against finite differences, like the other gradient tests, with a loss on the conics. With the previous backward pass, the same check is off by 11% to 95%, depending on the argument. With this change, it agrees within 5e-4. Both figures come from running the projection kernels on Warp's CPU device with the test's scene.
• On an RTX 4070 Ti, tests/unit/test_project.py, tests/unit/test_transforms.py, tests/integration/test_render_grad.py and tests/integration/test_transforms.py pass: 50 passed and 5 skipped. The skipped tests are the gsplat parity tests, because gsplat is not installed.
• ruff check and ruff format --check pass.

Impact

I render a flight hall (robot_hall.ply from amacati/splats) through Crazyflow's FPV camera, about 2 m above the floor at 160x120 px. Compared with finite differences, the gradient of an image loss with respect to the camera position was off by 66% along the optical axis, 24% sideways and 5% vertically. With this change, the errors are 2.3%, 2.3% and 1.9%.

[amacati]

Hi, thanks for the PR. I will hopefully take a look at it later today. Just out of curiosity, what are you using splax for?

[manuel-rodriguezvillegas]

Hi, thanks for the PR. I will hopefully take a look at it later today. Just out of curiosity, what are you using splax for?

Hello! We are working on vision-only FPV flight with differentiable physics and were impressed by the Crazyflow results, so we wanted to try it with 3DGS as an environment representation.

[amacati]

Left some suggestions to remove LLM-like wording in comments, simplify the test and expand coverage. Otherwise this looks good.

[manuel-rodriguezvillegas]

Thanks for the review! All suggestions applied. On an RTX 4070 Ti, test_project.py, test_transforms.py, test_render_grad.py and the integration test_transforms.py pass (51 passed, 5 skipped for gsplat parity), including both lens cases of test_project_grad_clamped.

[amacati]

Thanks for fixing this! Will try to make a minor release soon so that the fix is on PyPI.