ai-edge-torch-nightly 0.3.0.dev20240827__py3-none-any.whl → 0.3.0.dev20240829__py3-none-any.whl

ai_edge_torch/odml_torch/export_utils.py

@@ -0,0 +1,168 @@
+# 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 ODML Torch export."""
+
+import functools
+import re
+from typing import Sequence, cast
+import jax._src.interpreters.mlir
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import func
+import torch
+
+# std::numeric_limits<int64_t>::min()
+IR_DYNAMIC = -9223372036854775808
+
+
+def is_ir_dynamic(v):
+  return v == IR_DYNAMIC
+
+
+def is_torch_dynamic(v):
+  return isinstance(v, torch.SymInt)
+
+
+def is_iterable(v):
+  try:
+    iter(v)
+  except TypeError:
+    return False
+  return True
+
+
+def create_ir_context():
+  # HACK: Use ir context from JAX as base for better stability in OSS.
+  # TODO(b/362798610) Build MLIR pybinding in ai-edge-torch release.
+  context = jax._src.interpreters.mlir.make_ir_context()
+  context.allow_unregistered_dialects = True
+
+  return context
+
+
+def inline(
+    symbol_table: ir.SymbolTable,
+    block: ir.Block,
+):
+  """Recursively inlines all func.call ops in the block.
+
+  The symbol_table must include all func.func called by func.call ops.
+  This inliner in Python is implemented because MLIR inline pass from JAX's
+  MLIR pybinding build in OSS cannot properly inline func.call ops.
+  """
+  while True:
+    is_changed = False
+    for op in block.operations:
+      if op.OPERATION_NAME != func.CallOp.OPERATION_NAME:
+        continue
+
+      call_op = cast(func.CallOp, op)
+      func_op = cast(func.FuncOp, symbol_table[call_op.callee.value])
+      with ir.InsertionPoint(op):
+        new_results = clone_func_body_ops(func_op, call_op.operands)
+
+      for old_result, new_result in zip(call_op.results, new_results):
+        old_result = cast(ir.Value, old_result)
+        old_result.replace_all_uses_with(new_result)
+      call_op.erase()
+      is_changed = True
+
+    if not is_changed:
+      break
+
+  for op in block.operations:
+    for region in op.regions:
+      for block in region.blocks:
+        inline(symbol_table, block)
+
+
+def clone_func_body_ops(func_op: func.FuncOp, ir_inputs: Sequence[ir.Value]):
+  """Clone operations in the func_op's body by one into the current context."""
+  func_args = list(func_op.arguments)
+  ir_inputs = list(ir_inputs)
+  assert len(func_args) == len(ir_inputs)
+
+  value_mapping = {arg: ir_input for arg, ir_input in zip(func_args, ir_inputs)}
+
+  for op in list(func_op.entry_block.operations):
+    cloned_operands = [value_mapping[val] for val in op.operands]
+    if op.OPERATION_NAME == func.ReturnOp.OPERATION_NAME:
+      return cloned_operands
+
+    cloned = cast(ir.Operation, op.operation.clone())
+
+    for i in range(len(op.operands)):
+      cloned.operands[i] = cloned_operands[i]
+
+    for i in range(len(op.results)):
+      value_mapping[op.results[i]] = cloned.results[i]
+
+  return []
+
+
+def sanitize_aten_op_name(op, chars=":."):
+  return re.sub("[{}]".format(chars), "_", str(op))
+
+
+def build_ir_attr(val):
+  if val is None:
+    return ir.StringAttr.get("py_None")
+  if isinstance(val, bool):
+    return ir.BoolAttr.get(val)
+  if isinstance(val, int):
+    return ir.IntegerAttr.get(ir.IntegerType.get_signless(64), val)
+  if isinstance(val, float):
+    return ir.BoolAttr.get(val)
+  if isinstance(val, str):
+    return ir.StringAttr.get(val)
+  if isinstance(val, dict):
+    return ir.DictAttr.get({k: build_ir_attr(v) for k, v in val.items()})
+  if isinstance(val, (list, tuple)):
+    return ir.ArrayAttr.get([build_ir_attr(v) for v in val])
+
+  # Stringify the value to a StringAttr by default
+  return ir.StringAttr.get(str(val))
+
+
+def torch_dtype_to_ir_element_type(ctx, 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.bool: functools.partial(ir.IntegerType.get_signless, 1),
+  }.get(dtype)
+  return ty_get(ctx)
+
+
+def ir_element_type_to_torch_dtype(ty):
+  if isinstance(ty, ir.F32Type):
+    return torch.float32
+  if isinstance(ty, ir.F64Type):
+    return torch.float64
+  if isinstance(ty, ir.F16Type):
+    return torch.half
+  if isinstance(ty, ir.IntegerType):
+    if ty.is_signless:
+      if ty.width == 64:
+        return torch.long
+      if ty.width == 32:
+        return torch.int32
+      if ty.width == 16:
+        return torch.int16
+      if ty.width == 1:
+        return torch.bool
+  raise RuntimeError(f"Unsupported ir element type: {ty}")

ai_edge_torch/odml_torch/jax_bridge/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+# ==============================================================================
+from ai_edge_torch.odml_torch.jax_bridge._wrap import wrap

ai_edge_torch/odml_torch/jax_bridge/_wrap.py

@@ -0,0 +1,152 @@
+# 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.
+# ==============================================================================
+"""APIs to wrap JAX functions for using in ODML Torch lowerings."""
+
+import functools
+import inspect
+from typing import Any, Callable, cast
+import uuid
+from ai_edge_torch.odml_torch import export_utils
+from ai_edge_torch.odml_torch import passes
+from ai_edge_torch.odml_torch.jax_bridge import utils
+import jax
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import func
+import torch.utils._pytree as pytree
+
+# Jax double (64bit) precision is required to generate StableHLO mlir with
+# i64/f64 tensors from Jax bridged lowerings. If not set properly, all the
+# 64bit tensors would be truncated to 32bit dtype and potentially break the
+# lowering.
+jax.config.update("jax_enable_x64", True)
+
+
+def _lower_to_ir_text(
+    jaxfn, args, kwargs, ir_input_names: list[str] = None
+) -> str:
+  args = utils.tree_map_list_to_tuple(args)
+  kwargs = utils.tree_map_list_to_tuple(kwargs)
+
+  names_args = [
+      *zip(inspect.signature(jaxfn).parameters.keys(), args),
+      *kwargs.items(),
+  ]
+
+  static_argnames = []
+  jax_lower_static_kwargs = {}
+  jax_lower_args = []
+  jax_lower_argnames = []
+  ir_inputs = []
+
+  for i, (name, arg) in enumerate(names_args):
+    is_positional = i < len(args)
+    if not utils.is_ir_variable(arg):
+      static_argnames.append(name)
+      jax_lower_static_kwargs[name] = arg
+    else:
+      # Enforce the arg order in the mlir is the same as the lowering func
+      jax_lower_args.append(utils.ir_variable_to_jax(arg))
+
+      if is_positional and len(jax_lower_args) == i + 1:
+        # The first N continuous tensor args are passed to the lowering func
+        # as positional args, when they passed to the bridged func as
+        # positional args also.
+        jax_lower_argnames.append(None)
+      else:
+        # Otherwise pass the arg to the lowering func as keyword arg.
+        jax_lower_argnames.append(name)
+
+      if ir_input_names is None or name in ir_input_names:
+        # ir variable can be a nested tuple, while mlir args should be flat.
+        ir_inputs += [
+            x for x in pytree.tree_flatten(arg)[0] if isinstance(x, ir.Value)
+        ]
+
+  def new_lowering(*args, **jax_lower_static_kwargs):
+    jaxfn_args = []
+    jaxfn_kwargs = jax_lower_static_kwargs.copy()
+    for name, arg in zip(jax_lower_argnames, args):
+      if name is None:
+        jaxfn_args.append(arg)
+      else:
+        jaxfn_kwargs[name] = arg
+
+    return jaxfn(*jaxfn_args, **jaxfn_kwargs)
+
+  return (
+      jax.jit(new_lowering, static_argnames=static_argnames)
+      .lower(*jax_lower_args, **jax_lower_static_kwargs)
+      .as_text()
+  ), ir_inputs
+
+
+def wrap(jaxfn: Callable[Any, Any], ir_input_names: list[str] = None):
+  """Return the wrapped JAX function to be used in ODMLTorch lowerings.
+
+  If the given jaxfn has signature `jaxfn(*args, **kwargs) -> return`, the
+  wrapped function would:
+  - Have signature `wrapped(lctx: odml_torch.export.LoweringContext, *args,
+  **kwargs) -> return`.
+  - Accept mlir.ir.Value for all params expecting jax.Array as inputs.
+  - Return mlir.ir.Value for all jax.Array outputs from jaxfn.
+
+  Args:
+    jaxfn: The JAX function to be wrapped.
+    ir_input_names: The input (param) names of the JAX function to be used in
+      the MLIR lowering. This is useful when the JAX impl only depends on
+      specific inputs to the function. If not specified, all ir.Value passed to
+      the wrapped function are assumed to be used in the lowering.
+  """
+
+  @functools.wraps(jaxfn)
+  def wrapped(lctx, *args, **kwargs):
+
+    ir_text, ir_inputs = _lower_to_ir_text(
+        jaxfn,
+        args,
+        kwargs,
+        ir_input_names=ir_input_names,
+    )
+
+    module = ir.Module.parse(ir_text)
+    passes.strip_debuginfo(module)
+
+    symbol_table = ir.SymbolTable(module.operation)
+    main_func = symbol_table["main"]
+
+    with ir.InsertionPoint.at_block_begin(lctx.ir_module.body):
+      cloned_func = cast(func.FuncOp, main_func.clone())
+      cloned_func_name = f"{jaxfn.__name__}_{uuid.uuid4().hex[:8]}"
+      cloned_func.attributes["sym_name"] = ir.StringAttr.get(cloned_func_name)
+      cloned_func.attributes["sym_visibility"] = ir.StringAttr.get("private")
+
+    # HACK: Use the custom inliner implemented in Python because MLIR inline
+    # pass from JAX's MLIR pybinding build in OSS cannot properly inline
+    # func.call ops.
+    # This should be switched to `passes.inline(module)` when we have our own
+    # MLIR pybinding build.
+    export_utils.inline(symbol_table, cloned_func.entry_block)
+
+    if not cloned_func.arguments:
+      # Known edge case: when the lowering does not depend on input but
+      # just the meta of input like shape or dtype.
+      ir_inputs = []
+
+    results = func.CallOp(cloned_func, ir_inputs).results
+    if len(results) == 1:
+      return results[0]
+    return results
+
+  return wrapped

ai_edge_torch/odml_torch/jax_bridge/utils.py

@@ -0,0 +1,75 @@
+# 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 Jax bridge."""
+
+from ai_edge_torch import odml_torch
+import jax
+import jax.numpy as jnp
+from jax._src.lib.mlir import ir
+import torch
+
+
+def t2j_dtype(dtype):
+  return {
+      torch.bfloat16: jnp.bfloat16,
+      torch.half: jnp.float16,
+      torch.float32: jnp.float32,
+      torch.double: jnp.double,
+      torch.long: jnp.int64,
+      torch.int64: jnp.int64,
+      torch.int32: jnp.int32,
+      torch.int16: jnp.int16,
+      torch.int8: jnp.int8,
+      torch.uint8: jnp.uint8,
+      torch.bool: jnp.bool_,
+      torch.complex64: jnp.complex64,
+      torch.complex128: jnp.complex128,
+  }.get(dtype)
+
+
+def is_ir_variable(value):
+  if isinstance(value, ir.Value):
+    return True
+  if isinstance(value, (list, tuple)):
+    return any(is_ir_variable(x) for x in value)
+  return False
+
+
+def ir_variable_to_jax(value):
+  if isinstance(value, (list, tuple)):
+    return tuple([ir_variable_to_jax(x) for x in value])
+  elif not isinstance(value, ir.Value):
+    return value
+  elif not isinstance(value.type, ir.RankedTensorType):
+    raise ValueError(
+        f"ir.Value to JAX must be in ir.RankedTensorType, got {value}"
+    )
+
+  return jax.ShapeDtypeStruct(
+      value.type.shape,
+      t2j_dtype(
+          odml_torch.export_utils.ir_element_type_to_torch_dtype(
+              value.type.element_type
+          )
+      ),
+  )
+
+
+def tree_map_list_to_tuple(value):
+  if isinstance(value, dict):
+    return {k: tree_map_list_to_tuple(v) for k, v in value.items()}
+  if isinstance(value, (list, tuple)):
+    return tuple([tree_map_list_to_tuple(v) for v in value])
+  return value

ai_edge_torch/odml_torch/lowerings/__init__.py

@@ -0,0 +1,24 @@
+# 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.
+# ==============================================================================
+from . import _basic
+from . import _batch_norm
+from . import _convolution
+from . import _jax_lowerings
+from . import context
+from . import registry
+from . import utils
+from .registry import decompositions
+from .registry import lookup
+from .registry import lower

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -0,0 +1,204 @@
+# 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 math
+from typing import Optional, Union
+
+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 numpy as np
+import torch
+
+from .registry import lower
+
+
+# add(Tensor self, Tensor other) -> Tensor
+# @lower(torch.ops.aten.add)
+def _aten_add(lctx, x: ir.Value, y: ir.Value, alpha=1):
+  x, y = utils.upcast_to_same_type(x, y)
+  x, y = utils.broadcast_args_if_needed(x, y)
+  if alpha == 1:
+    return stablehlo.add(x, y)
+
+  alpha_splat = utils.splat(alpha, y.type.element_type, y.type.shape)
+  return stablehlo.add(x, stablehlo.multiply(y, alpha_splat))
+
+
+# mul.Tensor(Tensor self, Tensor other) -> Tensor
+# @lower(torch.ops.aten.mul.Tensor)
+def _aten_mul_tensor(lctx, self: ir.Value, other: ir.Value):
+  self, other = utils.upcast_to_same_type(self, other)
+  self, other = utils.broadcast_args_if_needed(self, other)
+
+  return stablehlo.multiply(self, other)
+
+
+# cat(Tensor[] tensors, int dim=0) -> Tensor
+# @lower(torch.ops.aten.cat)
+def _aten_cat(lctx, tensors: list[ir.Value], dim: int = 1):
+  return stablehlo.ConcatenateOp(tensors, dim).result
+
+
+# view(Tensor(a) self, SymInt[] size) -> Tensor(a)
+# @lower(torch.ops.aten.view)
+def _aten_view(lctx, self: ir.Value, size: list[int]):
+  return stablehlo.ReshapeOp(
+      ir.RankedTensorType.get(size, self.type.element_type), self
+  ).result
+
+
+# hardtanh(Tensor self, Scalar min_val=-1, Scalar max_val=1) -> Tensor
+@lower(torch.ops.aten.hardtanh)
+def _aten_hardtanh(
+    lctx,
+    self: ir.Value,
+    min_val: Union[int, float] = -1.0,
+    max_val: Union[int, float] = 1.0,
+):
+  elty = self.type.element_type
+  min_val = utils.splat(min_val, elty)
+  max_val = utils.splat(max_val, elty)
+
+  return stablehlo.clamp(min_val, self, max_val)
+
+
+# mean(Tensor self, *, ScalarType? dtype=None) -> Tensor
+# mean.dim(Tensor self, int[1]? dim, bool keepdim=False, *,
+#   ScalarType? dtype=None) -> Tensor
+@lower(torch.ops.aten.mean)
+@lower(torch.ops.aten.mean.dim)
+def _aten_mean_dim(
+    lctx,
+    self: ir.Value,
+    dim: Optional[list[int]] = None,
+    keepdim: bool = False,
+    *,
+    dtype=None,
+):
+  self_shape = self.type.shape
+  self_elty = self.type.element_type
+  if dim is None:
+    dim = list(range(len(self_shape)))
+  dim = [len(self_shape) + d if d < 0 else d for d in dim]
+  dim_ = ir.DenseI64ArrayAttr.get(np.asarray(dim, np.int64))
+  dim_to_keep = [d for d in range(len(self_shape)) if d not in dim]
+  dim_to_keep_ = ir.DenseI64ArrayAttr.get(np.asarray(dim_to_keep, np.int64))
+
+  zero_ = utils.splat(0.0, self_elty)
+
+  reduce_result_shape = [
+      s for d, s in enumerate(self_shape) if d in dim_to_keep
+  ]
+  reduce_result_ty = ir.RankedTensorType.get(reduce_result_shape, self_elty)
+  reduce_op = stablehlo.ReduceOp([reduce_result_ty], [self], [zero_], dim_)
+
+  reducer_arg_ty = ir.RankedTensorType.get(tuple(), self_elty)
+  reducer = reduce_op.regions[0].blocks.append(reducer_arg_ty, reducer_arg_ty)
+  with ir.InsertionPoint(reducer):
+    stablehlo.return_(
+        [stablehlo.add(reducer.arguments[0], reducer.arguments[1])]
+    )
+
+  sum_ = reduce_op.result
+  if keepdim:
+    sum_ = stablehlo.broadcast_in_dim(
+        ir.RankedTensorType.get(
+            [s if d in dim_to_keep else 1 for d, s in enumerate(self_shape)],
+            self_elty,
+        ),
+        sum_,
+        dim_to_keep_,
+    )
+
+  dim_els = math.prod([s for d, s in enumerate(self_shape) if d in dim])
+  dim_els_ = utils.splat(dim_els, self_elty)
+  div_ = stablehlo.broadcast_in_dim(
+      sum_.type, dim_els_, ir.DenseI64ArrayAttr.get([])
+  )
+  mean_ = stablehlo.divide(sum_, div_)
+
+  return mean_
+
+
+# https://pytorch.org/docs/stable/generated/torch.clone.html
+# https://github.com/pytorch/pytorch/blob/a95ceb51a23ae33c00b3a99224143c609b1b3eb3/aten/src/ATen/native/TensorFactories.cpp#L1730
+@lower(torch.ops.aten.clone)
+def _aten_clone(lctx, x: ir.Value, *, memory_format=None):
+  return x
+
+
+# https://pytorch.org/docs/stable/generated/torch.permute.html
+# https://github.com/pytorch/pytorch/blob/519151a062a9bd4f0d32a9c7c7eae47d7ed847b2/aten/src/ATen/native/TensorShape.cpp#L1448
+# https://github.com/openxla/stablehlo/blob/main/docs/spec.md#transpose
+@lower(torch.ops.aten.permute)
+def _aten_permute(lctx, x: ir.Value, dims: list[int]):
+  dim = len(x.type.shape)
+  return stablehlo.transpose(x, ir.DenseI64ArrayAttr.get(dims))
+
+
+# https://pytorch.org/docs/stable/generated/torch.mm.html
+# https://github.com/pytorch/pytorch/blob/ffabb25c489df1dc631a577c12a0c843c8b202f3/aten/src/ATen/native/LinearAlgebra.cpp#L193
+# https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dot_general
+@lower(torch.ops.aten.mm)
+def _aten_mm(mod, mat1: ir.Value, mat2: ir.Value) -> ir.Value:
+  mat1_shape = mat1.type.shape
+  mat2_shape = mat2.type.shape
+  mat1_dims = len(mat1_shape)
+  mat2_dims = len(mat2_shape)
+
+  if mat1_dims != 2 or mat1_dims != 2:
+    raise ValueError(
+        "Both arguments must be 2D matrices, received dimensions %d and %d"
+        % (mat1_dims, mat2_dims)
+    )
+
+  if mat1_shape[1] != mat2_shape[0]:
+    raise ValueError(
+        "mat1 and mat2 shapes cannot be multiplied, received shapes %s and %s"
+        % (mat1_shape, mat2_shape)
+    )
+
+  dot_dnums = stablehlo.DotDimensionNumbers.get(
+      lhs_batching_dimensions=[],
+      rhs_batching_dimensions=[],
+      lhs_contracting_dimensions=(1,),
+      rhs_contracting_dimensions=(0,),
+  )
+  return stablehlo.dot_general(
+      ir.RankedTensorType.get(
+          (mat1.type.shape[0], mat2.type.shape[1]), mat1.type.element_type
+      ),
+      mat1,
+      mat2,
+      dot_dnums,
+  )
+
+
+# https://pytorch.org/docs/stable/generated/torch.div.html
+# https://openxla.org/stablehlo/spec#divide
+# TODO: support rounding mode and type promotion (see torch.div spec).
+# @lower(torch.ops.aten.div)
+def _aten_div(mod, x, y, *, rounding_mode=None, out=None) -> ir.Value:
+  # By default, PyTorch performs a "true" division like Python 3. This requires
+  # casting integer input types to float to achieve the same semantics using
+  # stablehlo.divide.
+  if isinstance(x.type.element_type, ir.IntegerType):
+    x = utils.convert_int_to_float(x)
+  if isinstance(y.type.element_type, ir.IntegerType):
+    y = utils.convert_int_to_float(y)
+
+  x, y = utils.broadcast_args_if_needed(x, y)
+
+  return stablehlo.divide(x, y)