Challenges porting gaussian splatting to Jax

[peabody124]

Love Jax and Equinox, so want to use it for Gaussian Splatting. I've been trying to port the CUDA code for it over to Jax. I've got a lot of the pieces connected together, but am running into two key barriers. I suspect they can be overcome, but just not finding documentation around what I need:

• in the pytorch code it creates an empty torch::Tensor in the and passes a lambda to the CUDA code to resize. I can't find anything analogous to the dynamic memory allocation from Jax CUDA kernels. https://github.com/graphdeco-inria/diff-gaussian-rasterization/blob/main/rasterize_points.cu#L76
• these variables are stored in the runtime context for the backward pass. again, I'm sure this must be done for Jax code but cannot find any comparable examples https://github.com/graphdeco-inria/diff-gaussian-rasterization/blob/main/diff_gaussian_rasterization/__init__.py#L97

Any suggestions or pointers?

[jakevdp]

Thanks for the question! On the first point, if I understand correctly, JAX does not provide any equivalent API. All arrays in JAX must be statically-sized. In practice we've found that most use-cases that seem to require dynamic shapes can be re-expressed in terms of static shapes with some thought. Perhaps this can as well?

Regarding variables stored in the runtime context for the backward pass: this is part of the vector-jacobian product (VJP) code. For example, you can see the documentation of custom_vjp here: https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html#use-jax-custom-vjp-to-define-custom-reverse-mode-only-rules When you're doing reverse-mode differentiation over a sequence of standard operations in JAX, this runtime context is done implicitly. It's only when you're defining your own primitives or your own backward-pass logic via custom_vjp that you as the user need to think about this.

Does that help unblock you?

[peabody124]

Thanks! That reverse mode tip covers the storage of the intermediates.

I'll have to see if I can replicate this memory precomputing. It's rather tricky but should be deterministic. https://github.com/graphdeco-inria/diff-gaussian-rasterization/blob/main/cuda_rasterizer/rasterizer_impl.cu#L155

So as I understand, once I compute that, I should initialize that array in the lowering function and pass it to the custom call?

[yklcs]

@peabody124 Do you have any progress on this you can share? Very much interested in bringing Gaussian splatting to JAX.