Add dependency on ml_types

This will enable us to write bf16 numpy files.

Signed-off-by: Harsh Menon <[email protected]>

iree/turbine/kernel/wave/utils.py

@@ -59,6 +59,7 @@
 
 # TODO: Monkey-patching f16 support, need to fix in iree.
 import numpy
+import ml_dtypes
 
 bench.DTYPE_TO_ABI_TYPE[numpy.dtype(numpy.float16)] = "f16"
 
@@ -654,7 +655,16 @@ def compile_and_invoke(
             for inp in all_inputs:
                 if isinstance(inp, torch.Tensor):
                     tf = tempfile.NamedTemporaryFile(suffix=".npy")
-                    numpy.save(tf, inp.cpu().numpy())
+                    inp = inp.cpu()
+                    if inp.dtype == torch.bfloat16:
+                        inp = (
+                            inp.view(dtype=torch.uint16)
+                            .numpy()
+                            .view(dtype=ml_dtypes.bfloat16)
+                        )
+                    else:
+                        inp = inp.numpy()
+                    numpy.save(tf, inp)
                     tempfiles.append(tf)
                     inputs.append("@" + tf.name)
                 elif isinstance(inp, int):

requirements.txt

@@ -7,6 +7,7 @@ pytest==8.0.0
 pytest-xdist==3.5.0
 lit==18.1.7
 mypy==1.8.0
+ml_dtypes==0.5.0
 setuptools
 wheel