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

ai_edge_torch/odml_torch/jax_bridge/_wrap.py

@@ -0,0 +1,180 @@
+# 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
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+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
+) -> tuple[str, list[ir.Value]]:
+  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 lower_wrapper(*args):
+    nonlocal 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(lower_wrapper).lower(*jax_lower_args).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 lctx.node is None:
+      return results[0] if len(results) == 1 else results
+
+    out_avals = lctx.node.meta.get("tensor_meta") or lctx.node.meta.get("val")
+
+    if out_avals is None:
+      return results[0] if len(results) == 1 else results
+
+    def sanitize_result_elty(result, aval):
+      # JAX implementation may not respect aten op's output dtype. For example,
+      # JAX may implement a slightly different dtype upcast rules, leads to
+      # different result's dtype from bridged lowering and torch op output.
+      # Here we add an additional `stablehlo.convert` op when dtype does not
+      # match, to ensure the lowering's result dtype will always be the same
+      # as torch op's output dtype.
+      if aval is None:
+        return result
+
+      target_elty = export_utils.torch_dtype_to_ir_element_type(aval.dtype)
+      if result.type.element_type == target_elty:
+        return result
+      return stablehlo.convert(
+          ir.RankedTensorType.get(result.type.shape, target_elty), result
+      )
+
+    if len(results) == 1:
+      return sanitize_result_elty(results[0], out_avals)
+    return [
+        sanitize_result_elty(result, aval)
+        for result, aval in zip(results, out_avals)
+    ]
+
+  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,27 @@
+# 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 _layer_norm
+from . import _quantized_decomposed
+from . import _rand
+from . import context
+from . import registry
+from . import utils
+from .decomp import decompositions
+from .registry import lookup
+from .registry import lower

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -0,0 +1,294 @@
+# 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 import export_utils
+from ai_edge_torch.odml_torch.lowerings import context
+from ai_edge_torch.odml_torch.lowerings import registry
+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
+
+LoweringContext = context.LoweringContext
+lower = registry.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)
+
+
+# https://pytorch.org/docs/stable/generated/torch.floor.html
+# https://openxla.org/stablehlo/spec#floor
+@lower(torch.ops.aten.floor)
+def _aten_floor(lctx, x: ir.Value, *, out=None) -> ir.Value:
+  return stablehlo.floor(x)
+
+
+# Schema:
+#   - aten::cat(Tensor[] tensors, int dim=0) -> Tensor
+# Torch Reference:
+#   - https://pytorch.org/docs/main/generated/torch.cat.html
+@lower(torch.ops.aten.cat.default)
+def _aten_cat(lctx: LoweringContext, tensors, dim=0):
+  assert tensors
+  non_empty_tensors = [t for t in tensors if np.prod(t.type.shape) != 0]
+  out_aval = lctx.node.meta.get("tensor_meta") or lctx.node.meta.get("val")
+  if not non_empty_tensors:
+    return utils.splat(
+        0,
+        export_utils.torch_dtype_to_ir_element_type(out_aval.dtype),
+        out_aval.shape,
+    )
+
+  if dim < 0:
+    dim = dim + len(out_aval.shape)
+  dim = ir.IntegerAttr.get(ir.IntegerType.get_signless(64), dim)
+
+  return stablehlo.concatenate(non_empty_tensors, dim)
+
+
+# Schema:
+#   - aten::slice_scatter(Tensor self, Tensor src, int dim=0, SymInt?
+#       start=None, SymInt? end=None, SymInt step=1) -> Tensor
+# Torch Reference:
+#   - https://pytorch.org/docs/stable/generated/torch.slice_scatter.html
+#   - https://github.com/pytorch/pytorch/blob/18f9331e5deb4c02ae5c206e133a9b4add49bd97/aten/src/ATen/native/TensorShape.cpp#L4002
+@lower(torch.ops.aten.slice_scatter)
+def _aten_slice_scatter(lctx, self, src, dim=0, start=None, end=None, step=1):
+  start = start if start is not None else 0
+  end = end if end is not None else self.type.shape[dim]
+
+  start, end = np.clip(
+      [start, end], -self.type.shape[dim], self.type.shape[dim]
+  )
+
+  if start < 0:
+    start = self.type.shape[dim] + start
+  if end < 0:
+    end = self.type.shape[dim] + end
+
+  if end <= start or np.prod(src.type.shape) == 0:
+    return self
+
+  end = start + step * math.ceil((end - start) / step) - (step - 1)
+  padding_low = start
+  padding_high = self.type.shape[dim] - end
+  interior_padding = step - 1
+
+  rank = len(self.type.shape)
+  src = stablehlo.pad(
+      src,
+      utils.splat(0, src.type.element_type, []),
+      edge_padding_low=[padding_low if i == dim else 0 for i in range(rank)],
+      edge_padding_high=[padding_high if i == dim else 0 for i in range(rank)],
+      interior_padding=[
+          interior_padding if i == dim else 0 for i in range(rank)
+      ],
+  )
+
+  slices = [
+      slice(start, end, step) if i == dim else slice(None, None, None)
+      for i in range(rank)
+  ]
+  pred = np.ones(self.type.shape, dtype=np.bool_)
+  pred[np.index_exp[tuple(slices)]] = False
+  pred = stablehlo.constant(
+      ir.DenseElementsAttr.get(
+          np.packbits(pred, bitorder="little"),
+          type=ir.IntegerType.get_signless(1),
+          shape=pred.shape,
+      )
+  )
+  out = stablehlo.select(pred, self, src)
+  return out

ai_edge_torch/odml_torch/lowerings/_batch_norm.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.
+# ==============================================================================
+"""Provides lowering for coreaten to mlir stablehlo op: Convolution"""
+
+from typing import Optional
+
+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
+
+from .registry import lower
+
+
+# _native_batch_norm_legit_no_training(
+#     Tensor input,
+#     Tensor? weight,
+#     Tensor? bias,
+#     Tensor running_mean,
+#     Tensor running_var,
+#     float momentum,
+#     float eps) -> (Tensor, Tensor, Tensor)
+@lower(torch.ops.aten._native_batch_norm_legit_no_training)
+def _native_batch_norm_legit_no_training(
+    lctx,
+    input_tensor: ir.Value,
+    weight: Optional[ir.Value],
+    bias: Optional[ir.Value],
+    running_mean: ir.Value,
+    running_var: ir.Value,
+    momentum: float,
+    eps: float,
+):
+  if weight is None:
+    weight = utils.splat(
+        1, running_mean.type.element_type, running_mean.type.shape
+    )
+  if bias is None:
+    bias = utils.splat(
+        0, running_mean.type.element_type, running_mean.type.shape
+    )
+
+  return [
+      stablehlo.batch_norm_inference(
+          input_tensor, weight, bias, running_mean, running_var, eps, 1
+      ),
+      utils.splat(
+          0, input_tensor.type.element_type
+      ),  # TODO: return empty array instead
+      utils.splat(
+          0, input_tensor.type.element_type
+      ),  # TODO: return empty array instead
+  ]