host_callback to random number generator gives UserWarning: The value of the smallest subnormal ...

[Amith225]

Minimal Reproducible Code:

import jax.random
from jax import numpy as jnp, jit, ShapeDtypeStruct
from jax.experimental import host_callback


def keys(seed):
    key = jax.random.PRNGKey(seed)
    while True:
        key, subkey = jax.random.split(key)
        yield subkey


key_gen = keys(0)


def rng(shape):
    return jax.random.uniform(next(key_gen), shape)


@jit
def foo(x):
    rnd = host_callback.call(rng, x.shape, result_shape=ShapeDtypeStruct(x.shape, jnp.float32))
    return x * rnd


if __name__ == "__main__":
    print(foo(jnp.ones((3, 3))))

works all great but i get this weird warning:
UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero. setattr(self, word, getattr(machar, word).flat[0])

[jakevdp]

Thanks for the report – it looks like this somehow results from calling numpy.finfo in a host_callback context:

from jax.experimental import host_callback
from jax import jit, ShapeDtypeStruct
import numpy as np

def py_func(x):
  return np.finfo(x.dtype).min

@jit
def foo():
  return host_callback.call(py_func, 0.0, result_shape=ShapeDtypeStruct((), 'float32'))

foo()

/usr/local/lib/python3.9/dist-packages/numpy/core/getlimits.py:499: UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero.
  setattr(self, word, getattr(machar, word).flat[0])
/usr/local/lib/python3.9/dist-packages/numpy/core/getlimits.py:89: UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero.
  return self._float_to_str(self.smallest_subnormal)

I'm not sure why this would be the case; perhaps @sharadmv would have a guess

[hawkinsp]

I have a guess: XLA disables denormal values during XLA-generated code. When calling back into user-code from a callback, we're probably not re-enabling it.