[utils][torch] PyTorch to memref - enable bf16 conversion (#49)

Adds a special case for bfloat16 torch.Tensor to memref conversion.

Author: adam-smnk

lighthouse/utils/torch.py

@@ -1,4 +1,5 @@
 import ctypes
+import ml_dtypes
 
 import torch
 from mlir import ir
@@ -14,7 +15,15 @@ def to_memref(input: torch.Tensor) -> ctypes.Structure:
     Args:
         input: PyTorch tensor.
     """
-    return get_ranked_memref_descriptor(input.numpy())
+    if input.dtype == torch.bfloat16:
+        # Numpy doesn't support bf16 natively which disables
+        # direct conversion from PyTorch.
+        # Solved through non-destructive type casting.
+        nparray = input.view(dtype=torch.uint16).numpy()
+        nparray = nparray.view(ml_dtypes.bfloat16)
+    else:
+        nparray = input.numpy()
+    return get_ranked_memref_descriptor(nparray)
 
 
 def to_packed_args(inputs: list[torch.Tensor]) -> ctypes.Array[ctypes.c_void_p]:

pyproject.toml

@@ -17,7 +17,8 @@ dev = [
 
 [project.optional-dependencies]
 ingress_torch_mlir = [
-    "torch-mlir==20260125.703"
+    "torch-mlir==20260125.703",
+    "ml_dtypes",
 ]
 # Additional "targets" which pull in optional dependencies -- use `uv sync --extra TARGET`
 ingress_torch_cpu = [