ai-edge-torch-nightly 0.3.0.dev20250114__py3-none-any.whl

ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py

@@ -0,0 +1,177 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Lowerings for PT2E torch.ops.quantized_decomposed ops."""
+from typing import Optional, Union, cast
+
+from ai_edge_torch.odml_torch.lowerings import context
+from ai_edge_torch.odml_torch.lowerings import utils
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import torch
+import torch.ao.quantization.fx._decomposed
+import torch.utils._pytree as pytree
+
+from . import registry
+
+lower = registry.lower
+LoweringContext = context.LoweringContext
+
+
+def _uniform_quantized_type(
+    stored_type: Union[str, ir.Type],
+    expressed_type: Union[str, ir.Type],
+    *,
+    scale=Union[float, list[float], tuple[float]],
+    zero_point=Union[float, list[float], tuple[float]],
+    storage_type_min: Optional[int] = None,
+    storage_type_max: Optional[int] = None,
+    channel_axis: Optional[int] = None,
+    channel_axis_size: Optional[int] = None,
+):
+  """Polyfill for quant.UniformQuantizedType."""
+  if storage_type_min and storage_type_max:
+    storage_min_max = f"<{storage_type_min}:{storage_type_max}>"
+  else:
+    storage_min_max = ""
+
+  if channel_axis is not None:
+    # Per-channel quantization
+    # https://mlir.llvm.org/docs/Dialects/QuantDialect/#per-channel-quantization
+    assert isinstance(scale, (list, tuple))
+    assert isinstance(zero_point, (list, tuple))
+
+    scale = list(scale)
+    zero_point = list(zero_point)
+
+    if len(scale) == 1:
+      scale = scale * channel_axis_size
+    if len(zero_point) == 1:
+      zero_point = zero_point * channel_axis_size
+
+    assert len(scale) == len(zero_point) == channel_axis_size
+    scale_zp_strs = []
+    for s, zp in zip(scale, zero_point):
+      scale_zp_strs.append(f"{s}:{zp}")
+    scale_zp = "{" + ",".join(scale_zp_strs) + "}"
+    return ir.Type.parse(
+        f"!quant.uniform<{stored_type}{storage_min_max}:{expressed_type}:{channel_axis},{scale_zp}>"
+    )
+  else:
+    # Per-layer quantization
+    # https://mlir.llvm.org/docs/Dialects/QuantDialect/#per-layer-quantization
+    scale = pytree.tree_flatten([scale])[0][-1]
+    zero_point = pytree.tree_flatten([zero_point])[0][-1]
+    scale_zp = f"{scale}:{zero_point}"
+    return ir.Type.parse(
+        f"!quant.uniform<{stored_type}{storage_min_max}:{expressed_type},{scale_zp}>"
+    )
+
+
+# Quant dialect is not registered in the Python MLIR pybinding used by
+# odml-torch. Therefore, stablehlo.uniform_quantize/uniform_dequantize ops and
+# quant types are represented in stablehlo.custom_call to pass MLIR verification
+# and VHLO serialization before converter.
+# TODO(b/362798610) Build MLIR pybinding in ai-edge-torch release.
+
+
+# Schema:
+#   - quantized_decomposed::quantize_per_tensor(Tensor input, float scale,
+#       int zero_point, int quant_min, int quant_max,
+#       ScalarType dtype) -> Tensor
+#   - quantized_decomposed::quantize_per_tensor.tensor(Tensor input,
+#       Tensor scale, Tensor zero_point, int quant_min, int quant_max,
+#       ScalarType dtype) -> Tensor
+#
+# Scale and zero_point in tensors are automatically converted to list before
+# lowering.
+@lower(torch.ops.quantized_decomposed.quantize_per_tensor)
+def _quantize_per_tensor(
+    lctx: LoweringContext,
+    input: ir.Value,
+    scale: Union[float, list[float]],
+    zero_point: Union[float, list[float]],
+    quant_min: int,
+    quant_max: int,
+    dtype: torch.dtype,
+):
+  input_ty = cast(ir.RankedTensorType, input.type)
+  qty = _uniform_quantized_type(
+      utils.torch_dtype_to_ir_element_type(dtype),
+      input_ty.element_type,
+      scale=scale,
+      zero_point=zero_point,
+      storage_type_min=quant_min,
+      storage_type_max=quant_max,
+  )
+  return stablehlo.custom_call(
+      call_target_name="odml_torch.uniform_quantize",
+      inputs=[input],
+      result=[input_ty],
+      backend_config=ir.StringAttr.get(
+          str(ir.RankedTensorType.get(input_ty.shape, qty))
+      ),
+  )
+
+
+# Schema:
+#   - quantized_decomposed::quantize_per_channel(Tensor input, Tensor scales,
+#       Tensor zero_points, int axis, int quant_min, int quant_max,
+#       ScalarType dtype) -> Tensor
+#
+# Scale and zero_point in tensors are automatically converted to list before
+# lowering.
+@lower(torch.ops.quantized_decomposed.quantize_per_channel)
+def _quantize_per_channel(
+    lctx: LoweringContext,
+    input: ir.Value,
+    scale: list[float],
+    zero_point: list[float],
+    axis: int,
+    quant_min: int,
+    quant_max: int,
+    dtype: torch.dtype,
+):
+  input_ty = cast(ir.RankedTensorType, input.type)
+  qty = _uniform_quantized_type(
+      utils.torch_dtype_to_ir_element_type(dtype),
+      input_ty.element_type,
+      scale=scale,
+      zero_point=zero_point,
+      channel_axis=axis,
+      channel_axis_size=input_ty.shape[axis],
+      storage_type_min=quant_min,
+      storage_type_max=quant_max,
+  )
+  return stablehlo.custom_call(
+      call_target_name="odml_torch.uniform_quantize",
+      inputs=[input],
+      result=[input_ty],
+      backend_config=ir.StringAttr.get(
+          str(ir.RankedTensorType.get(input_ty.shape, qty))
+      ),
+  )
+
+
+@lower(torch.ops.quantized_decomposed.dequantize_per_tensor)
+@lower(torch.ops.quantized_decomposed.dequantize_per_channel)
+def _dequantize(lctx: LoweringContext, input: ir.Value, *args, **kwargs):
+  result_meta = lctx.node.meta.get("tensor_meta")
+  result_elty = utils.torch_dtype_to_ir_element_type(result_meta.dtype)
+
+  return stablehlo.custom_call(
+      call_target_name="odml_torch.uniform_dequantize",
+      inputs=[input],
+      result=[ir.RankedTensorType.get(result_meta.shape, result_elty)],
+  )

ai_edge_torch/odml_torch/lowerings/_rand.py

@@ -0,0 +1,142 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import uuid
+
+from ai_edge_torch.odml_torch import export_utils
+from ai_edge_torch.odml_torch.lowerings import context
+from ai_edge_torch.odml_torch.lowerings import registry
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import func
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import numpy as np
+import torch
+import torch.utils._pytree as pytree
+
+LoweringContext = context.LoweringContext
+lower = registry.lower
+
+
+def _random_lowering(
+    lctx: LoweringContext,
+    size: list[int],
+    generator,
+    dtype: torch.dtype,
+    rand_tensor,
+    composite_name: str,
+):
+  if dtype is None:
+    dtype = torch.float32
+
+  rand_tensor = rand_tensor.type(dtype)
+  data = rand_tensor.detach().numpy()
+
+  shape, _ = pytree.tree_flatten(size)
+  elty = export_utils.torch_dtype_to_ir_element_type(dtype)
+
+  decomp_name = f"{composite_name}.impl_{uuid.uuid4().hex[:8]}"
+
+  with ir.InsertionPoint(lctx.ir_module.body):
+
+    @func.FuncOp.from_py_func(
+        ir.RankedTensorType.get(
+            [len(shape)],
+            ir.IntegerType.get_signless(32),
+        ),
+        name=decomp_name,
+    )
+    def _rand_impl(_):
+      return [stablehlo.constant(ir.DenseElementsAttr.get(data))]
+
+  seed, seed2 = (
+      torch.randint(
+          torch.iinfo(torch.int64).min,
+          torch.iinfo(torch.int64).max,
+          (2,),
+          dtype=torch.int64,
+          generator=generator,
+      )
+      .detach()
+      .numpy()
+  )
+
+  shape_ = stablehlo.constant(
+      ir.DenseElementsAttr.get(np.array(shape, dtype=np.int32))
+  )
+  return stablehlo.CompositeOp(
+      result=[ir.RankedTensorType.get(shape, elty)],
+      inputs=[shape_],
+      name=composite_name,
+      composite_attributes=ir.DictAttr.get({
+          "seed": ir.IntegerAttr.get(ir.IntegerType.get_signless(64), seed),
+          "seed2": ir.IntegerAttr.get(ir.IntegerType.get_signless(64), seed2),
+      }),
+      decomposition=decomp_name,
+  ).results[0]
+
+
+# Schema:
+# - aten::rand(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None,
+#     Device? device=None, bool? pin_memory=None) -> Tensor
+# - aten::rand.generator(SymInt[] size, *, Generator? generator,
+#     ScalarType? dtype=None, Layout? layout=None, Device? device=None,
+#     bool? pin_memory=None) -> Tensor
+@registry.lower(torch.ops.aten.rand)
+def _aten_rand(
+    lctx: LoweringContext,
+    size,
+    generator=None,
+    dtype=None,
+    layout=torch.strided,
+    device=None,
+    pin_memory=False,
+):
+  return _random_lowering(
+      lctx,
+      size,
+      generator,
+      dtype,
+      rand_tensor=torch.ops.aten.rand.generator(
+          size, generator=generator, dtype=dtype
+      ),
+      composite_name="odml.random_uniform",
+  )
+
+
+# Schema:
+# - aten::randn(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None,
+#     Device? device=None, bool? pin_memory=None) -> Tensor
+# - aten::randn.generator(SymInt[] size, *, Generator? generator,
+#     ScalarType? dtype=None, Layout? layout=None, Device? device=None,
+#     bool? pin_memory=None) -> Tensor
+@registry.lower(torch.ops.aten.randn)
+def _aten_randn(
+    lctx: LoweringContext,
+    size,
+    generator=None,
+    dtype=None,
+    layout=torch.strided,
+    device=None,
+    pin_memory=False,
+):
+  return _random_lowering(
+      lctx,
+      size,
+      generator,
+      dtype,
+      rand_tensor=torch.ops.aten.randn.generator(
+          size, generator=generator, dtype=dtype
+      ),
+      composite_name="odml.random_standard_normal",
+  )

ai_edge_torch/odml_torch/lowerings/context.py

@@ -0,0 +1,42 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Define context object for export and MLIR lowerings."""
+
+import dataclasses
+from jax._src.lib.mlir import ir
+import torch
+
+
+@dataclasses.dataclass
+class LoweringContext:
+  """The context object used in export interpreter and MLIR lowerings."""
+
+  ir_context: ir.Context
+  ir_module: ir.Module
+  ir_location: ir.Location = None
+  node: torch.fx.Node = None
+
+  @property
+  def ctx(self):
+    """Shortcut for ir_context."""
+    return self.ir_context
+
+  @property
+  def loc(self):
+    """Shortcut for ir_location."""
+    return self.ir_location
+
+  def replace(self, **kwargs):
+    return dataclasses.replace(self, **kwargs)

ai_edge_torch/odml_torch/lowerings/decomp.py

@@ -0,0 +1,69 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Torch export decompositions to run before lowering."""
+
+import functools
+
+import torch
+
+
+@functools.cache
+def decompositions():
+  # Base: Core ATen decompositions
+  decompositions = torch._decomp.core_aten_decompositions()
+
+  decompositions.update(
+      torch._decomp.get_decompositions([
+          torch.ops.aten.upsample_nearest2d,
+          torch.ops.aten._native_batch_norm_legit.no_stats,
+          torch.ops.aten._native_batch_norm_legit_functional,
+          torch.ops.aten._adaptive_avg_pool2d,
+          torch.ops.aten._adaptive_avg_pool3d,
+          torch.ops.aten.grid_sampler_2d,
+          torch.ops.aten.native_group_norm,
+          torch.ops.aten.native_dropout,
+          torch.ops.aten.reflection_pad1d,
+          torch.ops.aten.reflection_pad2d,
+          torch.ops.aten.reflection_pad3d,
+          torch.ops.aten.replication_pad1d,
+          torch.ops.aten.replication_pad2d,
+          torch.ops.aten.replication_pad3d,
+          torch.ops.aten.addmm,
+      ])
+  )
+
+  torch._decomp.remove_decompositions(
+      decompositions,
+      [
+          torch.ops.aten.roll,
+          # Torch's default einsum impl/decompositions is less efficient and
+          # optimized through converter than JAX's impl. Disable einsum
+          # decomposition to use JAX bridge for a more efficient lowering.
+          torch.ops.aten.einsum.default,
+      ],
+  )
+
+  # Override noop aten op decompositions for faster run_decompositions.
+  decompositions[torch.ops.aten.alias.default] = lambda x: x
+  decompositions[torch.ops.aten.detach.default] = lambda x: x
+
+  # Override _safe_softmax decompositions with regular softmax.
+  # _safe_softmax introduces additional check-select ops to guard extreme
+  # input values to softmax, which could make the converted model inefficient
+  # on-device.
+  if hasattr(torch.ops.aten, "_safe_softmax"):
+    decompositions[torch.ops.aten._safe_softmax.default] = torch.softmax
+
+  return decompositions

ai_edge_torch/odml_torch/lowerings/registry.py

@@ -0,0 +1,65 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Torch op decompositions and MLIR lowerings registry."""
+
+from typing import Any, Callable
+
+import torch
+
+from . import context
+
+
+class LoweringRegistry:
+  """Registry object for torch op decompositions and to-MLIR lowerings."""
+
+  def __init__(self):
+    self.registered_ops = {}
+
+  def lookup(self, op_or_name):
+    candidate = self._get_lowering(op_or_name)
+    if candidate is None:
+      if isinstance(op_or_name, torch._ops.OpOverloadPacket):
+        candidate = self._get_lowering(op_or_name.default)
+      if isinstance(op_or_name, torch._ops.OpOverload):
+        candidate = self._get_lowering(op_or_name.overloadpacket)
+    return candidate
+
+  def _get_lowering(self, op):
+    candidate = self.registered_ops.get(op)
+    return candidate
+
+  def register(self, op, lowering):
+    if isinstance(op, torch._ops.OpOverloadPacket):
+      ops = [getattr(op, overload) for overload in op.overloads()]
+    else:
+      ops = [op]
+
+    for op in ops:
+      self.registered_ops[op] = lowering
+
+
+global_registry = LoweringRegistry()
+
+
+def lookup(op):
+  return global_registry.lookup(op)
+
+
+def lower(op):
+  def inner(lowering: Callable[[context.LoweringContext, ...], Any]):
+    global_registry.register(op, lowering)
+    return lowering
+
+  return inner

ai_edge_torch/odml_torch/lowerings/utils.py

@@ -0,0 +1,201 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for building MLIR lowerings."""
+
+import functools
+import numbers
+from typing import Any
+from typing import Optional
+
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import numpy as np
+import torch
+
+
+def torch_dtype_to_ir_element_type(dtype):
+  ty_get = {
+      torch.double: ir.F64Type.get,
+      torch.float32: ir.F32Type.get,
+      torch.half: ir.F16Type.get,
+      torch.long: functools.partial(ir.IntegerType.get_signless, 64),
+      torch.int32: functools.partial(ir.IntegerType.get_signless, 32),
+      torch.int16: functools.partial(ir.IntegerType.get_signless, 16),
+      torch.int8: functools.partial(ir.IntegerType.get_signless, 8),
+      torch.bool: functools.partial(ir.IntegerType.get_signless, 1),
+  }[dtype]
+  return ty_get()
+
+
+def splat(val, ty, shape=tuple(), *, loc: Optional[Any] = None):
+  if isinstance(ty, ir.IntegerType):
+    if ty.width == 1:
+      attr = ir.BoolAttr.get(bool(val))
+    else:
+      attr = ir.IntegerAttr.get(ty, int(val))
+  elif isinstance(ty, ir.FloatType):
+    attr = ir.FloatAttr.get(ty, val)
+  else:
+    raise ValueError("Unsupported type: %s" % str(ty))
+
+  return stablehlo.constant(
+      ir.DenseElementsAttr.get_splat(
+          ir.RankedTensorType.get(shape, ty),
+          attr,
+      ),
+      loc=loc,
+  )
+
+
+def get_common_broadcast_shape(
+    shape_1: list[int], shape_2: list[int]
+) -> Optional[list[int]]:
+  if not shape_1 and not shape_2:
+    return None
+
+  shape_1 = shape_1 if shape_1 else [1]
+  shape_2 = shape_2 if shape_2 else [1]
+
+  length_diff = abs(len(shape_1) - len(shape_2))
+  if len(shape_1) < len(shape_2):
+    shape_1 = [1] * length_diff + shape_1
+  elif len(shape_1) > len(shape_2):
+    shape_2 = [1] * length_diff + shape_2
+
+  common_broadcast_shape = []
+  for idx in reversed(range(len(shape_1))):
+    dim_size1 = shape_1[idx]
+    dim_size2 = shape_2[idx]
+
+    if dim_size1 == dim_size2:
+      common_broadcast_shape.insert(0, dim_size1)
+    elif dim_size1 == 1 or dim_size2 == 1:
+      common_broadcast_shape.insert(0, max(dim_size1, dim_size2))
+    else:
+      return None
+
+  return common_broadcast_shape
+
+
+def get_broadcast_dimensions(
+    shape_from: list[int], shape_to: list[int]
+) -> list[int]:
+  assert get_common_broadcast_shape(shape_from, shape_to) == shape_to
+
+  ret = []
+  for val in range(len(shape_to) - len(shape_from), len(shape_to)):
+    ret.append(val)
+
+  return ir.DenseI64ArrayAttr.get(np.asarray(ret, np.int64))
+
+
+def broadcast_args_if_needed(
+    val_1: ir.Value, val_2: ir.Value
+) -> tuple[Optional[ir.Value], Optional[ir.Value]]:
+  broadcast_shape = get_common_broadcast_shape(
+      val_1.type.shape, val_2.type.shape
+  )
+  if broadcast_shape is None:
+    return None, None
+
+  new_val_1, new_val_2 = val_1, val_2
+
+  if val_1.type.shape != broadcast_shape:
+    new_val_1 = stablehlo.broadcast_in_dim(
+        result=ir.RankedTensorType.get(
+            broadcast_shape, val_1.type.element_type
+        ),
+        operand=val_1,
+        broadcast_dimensions=get_broadcast_dimensions(
+            val_1.type.shape, broadcast_shape
+        ),
+    )
+  if val_2.type.shape != broadcast_shape:
+    new_val_2 = stablehlo.broadcast_in_dim(
+        result=ir.RankedTensorType.get(
+            broadcast_shape, val_2.type.element_type
+        ),
+        operand=val_2,
+        broadcast_dimensions=get_broadcast_dimensions(
+            val_2.type.shape, broadcast_shape
+        ),
+    )
+  return new_val_1, new_val_2
+
+
+def upcast_to_same_type(*vals: ir.Value):
+  if not vals:
+    return None
+  if len(vals) == 1:
+    return vals[0]
+
+  def get_priority(ty: ir.Type):
+    priorities = [
+        ir.IntegerType.get_signless(1),
+        ir.IntegerType.get_signless(16),
+        ir.IntegerType.get_signless(32),
+        ir.IntegerType.get_signless(64),
+        ir.F16Type,
+        ir.F32Type,
+        ir.F64Type,
+    ]
+    for i, tycls in enumerate(priorities):
+      if tycls.isinstance(ty):
+        return i
+    raise ValueError("Unsupported type: %s" % str(ty))
+
+  cast_tycls = type(max([v.type.element_type for v in vals], key=get_priority))
+  new_vals = []
+  for val in vals:
+    if not cast_tycls.isinstance(val.type.element_type):
+      val = stablehlo.convert(
+          ir.RankedTensorType.get(val.type.shape, cast_tycls.get()), val
+      )
+    new_vals.append(val)
+  return tuple(new_vals)
+
+
+def minmax(ty: ir.Type) -> tuple[numbers.Number, numbers.Number]:
+  if isinstance(ty, ir.IntegerType):
+    if ty.is_unsigned:
+      return (0, 1 << ty.width)
+    else:
+      return (-(1 << (ty.width - 1)), (1 << (ty.width - 1)) - 1)
+  elif isinstance(ty, ir.F16Type):
+    return (np.finfo(np.float16).min, np.finfo(np.float16).max)
+  elif isinstance(ty, ir.F32Type):
+    return (np.finfo(np.float32).min, np.finfo(np.float32).max)
+  elif isinstance(ty, ir.F64Type):
+    return (np.finfo(np.float64).min, np.finfo(np.float64).max)
+  else:
+    raise ValueError("Unsupported type: %s" % ty)
+
+
+def convert_int_to_float(t: ir.Value) -> ir.Value:
+  """Converts an input with type ir.IntegerType to an ir.FloatType of equivalent width."""
+  elty = t.type.element_type
+  if not isinstance(elty, ir.IntegerType):
+    raise ValueError(
+        "Expected input with integer type, received %s" % type(elty)
+    )
+
+  if elty.width == 32:
+    return stablehlo.convert(
+        ir.RankedTensorType.get(t.type.shape, ir.F32Type.get()), t
+    )
+  elif elty.width == 64:
+    return stablehlo.convert(
+        ir.RankedTensorType.get(t.type.shape, ir.F64Type.get()), t
+    )