where is shjax ?

[Chutlhu]

Dear author,

thank you very much for this repository. I am interested in the spherical harmonics hash encoding, could you provide some more information about the shjax library? I cannot find it online

thank you very much

[blurgyy]

Hi @Chutlhu,

Thank you for your interest.

I implemented shjax as a custom extension in deps/spherical-harmonics-encoding-jax/, and it is integrated in python via

https://github.com/blurgyy/jaxngp/blob/d63c2c9f30b5e77c0ea212ceab94d62529f7a887/models/encoders.py#L351-L357

You can inspect its source there.

However, I ended up using a JAX implementation of the spherical harmonics encoding, because in my benchmarks, the JAX implementation is consistently faster than the custom CUDA implementation, I think it is because JAX code can be easier optimized via operations like kernel fusion. The JAX implementation which is used throughout the project can be found at https://github.com/blurgyy/jaxngp/blob/d63c2c9f30b5e77c0ea212ceab94d62529f7a887/models/encoders.py#L361

The benchmark I used to compare between the JAX vs CUDA implementation of spherical harmonics encoding is at https://github.com/blurgyy/jaxngp/blob/d63c2c9f30b5e77c0ea212ceab94d62529f7a887/models/encoders.py#L467

Cheers!

[Chutlhu]

Dear @blurgyy , Thank you very much! I found it. I played a little bit with it. Using this positional embedding, the model overfits the training data very well (compared to standard Random Fourier Features), but it seems to lose the native interpolation property. Do you know something about it? Do you have some references about this problem?