jax.random + ml_dtypes low memory types

[AlejandroBaron]

Hey I'm trying to use int4 and float8 from https://github.com/jax-ml/ml_dtypes. However there seems to be no support

from ml_dtypes import int4, float8_e5m2
from jax import random
key = random.PRNGKey(0)
random.uniform(key, (1,1), dtype=float8_e5m2)
random.randint(key, (1,1), 0, 10, dtype=int4)

Is this expected to be implemented in the near future?

[jakevdp]

We have generally been implementing quantized operations when there is (1) good motivation to do so, and (2) fast hardware to back it up. Can you say more about your use-case for these operations?

[kmheckel]

@jakevdp I am facing this same issue where I am trying to use float8 datatypes as part of training quantized spiking neural networks but cannot initialise weights for a network because jax.random._uniform does not support jnp.float8_* dtypes. My intent is to train on an H100 GPU which supports fp8 with double the throughput of fp16. The error message for jax.random._uniform seems to be a little out of date as well as it only discusses fp32/64 while also accepting fp16.

I'm not sure how much work it would take to permit fp8 types but would greatly appreciate it! It seems that the uniform sampling method would need updated but I'm not sure how many other changes would be required to support fp8. Beyond language model applications, neuroevolution would largely benefit from fp8 because of the fact that you could double the population size for the same compute/memory capacity and SNN research would also benefit as one could represent both the input spikes and model weights in fp8 and eliminate having to recast int8 inputs to fp16 values for use.

[jakevdp]

(continued in #16673)