[JAX] Why no higher precision?

[Qazalbash]

Hi, I wanted to know if there is any specific reason not to allow higher precision values?

Reference code:

gmmx/gmmx/gmm.py

Lines 97 to 99 in 537cfdd

	if array.dtype != jnp.float32:
	message = f"Expected float32, got {array.dtype}"
	raise ValueError(message)

[adonath]

Thanks for the issue!

Hi, I wanted to know if there is any specific reason not to allow higher precision values?

Absolutely not, this must be some left over code from when I did debugging on the GPU. I can remove this and add a test for float64 precision.