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

jaxngp/models/encoders.py

Lines 351 to 357 in d63c2c9

	class SphericalHarmonicsEncoderCuda(Encoder):
	# highest degree
	L: int
	def __call__(self, dirs: jax.Array) -> jax.Array:
	"Just a thin wrapper on top of :func:`shjax.spherical_harmonics_encoding()`"
	return shjax.spherical_harmonics_encoding(dirs, self.L)

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

jaxngp/models/encoders.py

Line 361 in d63c2c9

class SphericalHarmonicsEncoder(Encoder):

The benchmark I used to compare between the JAX vs CUDA implementation of spherical harmonics encoding is at

jaxngp/models/encoders.py

Line 467 in d63c2c9

def bench_sh():

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?