TypeError: ufunc 'isnan' not supported for the input types

[carlthome]

NumPy throws when evaluating _float8_e8m0fnu_finfo, with a TypeError: ufunc 'isnan' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe'

Background

Stumbled into this when using TensorFlow Datasets (its CLI tfds) on Google Cloud with ml_dtypes==0.5.1 and numpy==1.23.5

Possibly that NumPy 2 has to be used today? That's been hard to upgrade to with Apache Beam however due to ABI changes but might be a path forward.

Logs

Traceback (most recent call last):
  File "/usr/local/bin/tfds", line 5, in <module>
    from tensorflow_datasets.scripts.cli.main import launch_cli
  File "/usr/local/lib/python3.10/site-packages/tensorflow_datasets/scripts/cli/main.py", line 37, in <module>
    from tensorflow_datasets.scripts.cli import convert_format
  File "/usr/local/lib/python3.10/site-packages/tensorflow_datasets/scripts/cli/convert_format.py", line 33, in <module>
    from tensorflow_datasets.scripts.cli import convert_format_utils
  File "/usr/local/lib/python3.10/site-packages/tensorflow_datasets/scripts/cli/convert_format_utils.py", line 51, in <module>
    ConvertFn = Callable[[tf.train.Example], bytes]
  File "/usr/local/lib/python3.10/site-packages/etils/epy/lazy_imports_utils.py", line 118, in __getattr__
    return getattr(self._module, name)
  File "/usr/local/lib/python3.10/functools.py", line 981, in __get__
    val = self.func(instance)
  File "/usr/local/lib/python3.10/site-packages/etils/epy/lazy_imports_utils.py", line 79, in _module
    module = importlib.import_module(self.module_name)
  File "/usr/local/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "/usr/local/lib/python3.10/site-packages/tensorflow/__init__.py", line 51, in <module>
    from ._api.v2 import compat
  File "/usr/local/lib/python3.10/site-packages/tensorflow/_api/v2/compat/__init__.py", line 37, in <module>
    from . import v1
  File "/usr/local/lib/python3.10/site-packages/tensorflow/_api/v2/compat/v1/__init__.py", line 31, in <module>
    from . import compat
  File "/usr/local/lib/python3.10/site-packages/tensorflow/_api/v2/compat/v1/compat/__init__.py", line 37, in <module>
    from . import v1
  File "/usr/local/lib/python3.10/site-packages/tensorflow/_api/v2/compat/v1/compat/v1/__init__.py", line 48, in <module>
    from tensorflow._api.v2.compat.v1 import lite
  File "/usr/local/lib/python3.10/site-packages/tensorflow/_api/v2/compat/v1/lite/__init__.py", line 9, in <module>
    from . import experimental
  File "/usr/local/lib/python3.10/site-packages/tensorflow/_api/v2/compat/v1/lite/experimental/__init__.py", line 8, in <module>
    from . import authoring
  File "/usr/local/lib/python3.10/site-packages/tensorflow/_api/v2/compat/v1/lite/experimental/authoring/__init__.py", line 8, in <module>
    from tensorflow.lite.python.authoring.authoring import compatible
  File "/usr/local/lib/python3.10/site-packages/tensorflow/lite/python/authoring/authoring.py", line 43, in <module>
    from tensorflow.lite.python import convert
  File "/usr/local/lib/python3.10/site-packages/tensorflow/lite/python/convert.py", line 27, in <module>
    from tensorflow.lite.python import util
  File "/usr/local/lib/python3.10/site-packages/tensorflow/lite/python/util.py", line 52, in <module>
    from jax import xla_computation as _xla_computation
  File "/usr/local/lib/python3.10/site-packages/jax/__init__.py", line 25, in <module>
    from jax._src.cloud_tpu_init import cloud_tpu_init as _cloud_tpu_init
  File "/usr/local/lib/python3.10/site-packages/jax/_src/cloud_tpu_init.py", line 17, in <module>
    from jax._src import config
  File "/usr/local/lib/python3.10/site-packages/jax/_src/config.py", line 27, in <module>
    from jax._src import lib
  File "/usr/local/lib/python3.10/site-packages/jax/_src/lib/__init__.py", line 87, in <module>
    import jaxlib.xla_client as xla_client
  File "/usr/local/lib/python3.10/site-packages/jaxlib/xla_client.py", line 30, in <module>
    import ml_dtypes
  File "/usr/local/lib/python3.10/site-packages/ml_dtypes/__init__.py", line 40, in <module>
    from ml_dtypes._finfo import finfo
  File "/usr/local/lib/python3.10/site-packages/ml_dtypes/_finfo.py", line 153, in <module>
    class finfo(np.finfo):  # pylint: disable=invalid-name,missing-class-docstring
  File "/usr/local/lib/python3.10/site-packages/ml_dtypes/_finfo.py", line 699, in finfo
    _finfo_cache = {
  File "/usr/local/lib/python3.10/site-packages/ml_dtypes/_finfo.py", line 700, in <dictcomp>
    t: init_fn.__func__() for t, init_fn in _finfo_type_map.items()  # pytype: disable=attribute-error
  File "/usr/local/lib/python3.10/site-packages/ml_dtypes/_finfo.py", line 668, in _float8_e8m0fnu_finfo
    if not hasattr(obj, "tiny"):
  File "/usr/local/lib/python3.10/site-packages/numpy/core/getlimits.py", line 578, in tiny
    return self.smallest_normal
  File "/usr/local/lib/python3.10/site-packages/numpy/core/getlimits.py", line 557, in smallest_normal
    if isnan(self._machar.smallest_normal.flat[0]):
TypeError: ufunc 'isnan' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe''

[jakevdp]

This definitely works with more recent NumPy versions:

In [1]: import ml_dtypes

In [2]: ml_dtypes.__version__
Out[2]: '0.5.1'

In [3]: import numpy as np

In [4]: np.__version__
Out[4]: '2.2.2'

In [5]: x = np.arange(4, dtype=ml_dtypes.float8_e3m4)

In [6]: np.isnan(x)
Out[6]: array([False, False, False, False])

I'm not sure what exactly is different in older NumPy – perhaps @hawkinsp knows?

[hawkinsp]

Even if we could figure this out, the fix might well be "upgrade numpy", given it apparently works in newer NumPy. I suspect the problem is that NumPy doesn't really know that ml_dtypes types are floating point types, and we don't always perfectly hide that from the user.