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

ai_edge_torch/odml_torch/lowerings/_convolution.py

@@ -0,0 +1,243 @@
+# 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 stablehlo for Convolution."""
+
+import math
+from typing import Optional
+
+from ai_edge_torch.odml_torch.lowerings import registry
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import torch
+
+
+def make_padding(padding):
+  """Change the padding from pytorch to stablehlo style.
+
+  Stablehlo allows start and end padding for each dimension while aten only
+  allows symmetric padding and so only has one number per dimension.
+
+  Args:
+    padding: The padding of the convolution
+
+  Returns:
+    The padding in stablehlo style
+  """
+  return tuple((p, p) for p in padding)
+
+
+def create_conv_dimension_numbers(lhs, transposed: bool = False):
+  """Create the dimension numbers for the convolution.
+
+  Args:
+    lhs: The input tensor
+    transposed: Whether the convolution is transposed
+
+  Returns:
+    The dimension numbers for the convolution
+  """
+  num_spatial_dims = len(lhs.type.shape) - 2
+  spatial_dimensions = []
+  for i in range(0, num_spatial_dims):
+    spatial_dimensions.append(i + 2)
+
+  # Regular kernels are OIHW
+  # TransposedConv kernels are IOHW
+  dimension_numbers = stablehlo.ConvDimensionNumbers.get(
+      input_batch_dimension=0,
+      input_feature_dimension=1,
+      input_spatial_dimensions=spatial_dimensions,
+      kernel_input_feature_dimension=0 if transposed else 1,
+      kernel_output_feature_dimension=1 if transposed else 0,
+      kernel_spatial_dimensions=spatial_dimensions,
+      output_batch_dimension=0,
+      output_feature_dimension=1,
+      output_spatial_dimensions=spatial_dimensions,
+  )
+  return dimension_numbers
+
+
+def infer_output_shape(
+    lhs,
+    rhs,
+    stride,
+    dilation,
+    padding,
+    transposed: bool = False,
+    output_padding: list[int] = 0,
+):
+  """Infer the output shape of the convolution.
+
+  Args:
+    lhs: The input tensor
+    rhs: The kernel tensor
+    stride: The stride of the convolution (dilation of input in transposed conv)
+    dilation: The kernel dilation of the convolution
+    padding: The padding of the convolution
+    transposed: Whether the convolution is transposed
+    output_padding: The output padding of the convolution
+
+  Returns:
+    The output shape of the convolution
+  """
+  lhs_type: ir.RankedTensorType = lhs.type
+  lhs_shape: list[int] = lhs_type.shape
+  rhs_shape: list[int] = rhs.type.shape
+
+  # Input layout is: (N)CHW and Kernel layout is: (O)IHW for regular conv
+  # Input layout is: (N)CHW and Kernel layout is: I(O)HW for transposed conv
+  output_shape = (
+      [lhs_shape[0], rhs_shape[1]]
+      if transposed
+      else [lhs_shape[0], rhs_shape[0]]
+  )
+  num_spatial_dims = len(lhs.type.shape) - 2
+
+  # looping over the spatial dims (skipping the first 2 dims which are
+  # batch and features)
+  for spatial_dim in range(0, num_spatial_dims):
+    dim = spatial_dim + 2
+    dim_size = lhs_shape[dim]
+    kernel_dim_size = rhs_shape[dim]
+
+    if transposed:
+      output_dim_size = (
+          (dim_size - 1) * stride[spatial_dim]
+          - 2 * padding[spatial_dim]
+          + dilation[spatial_dim] * (kernel_dim_size - 1)
+          + output_padding[spatial_dim]
+          + 1
+      )
+    else:
+      output_dim_size = math.floor(
+          (
+              (
+                  dim_size
+                  + 2 * padding[spatial_dim]
+                  - dilation[spatial_dim] * (kernel_dim_size - 1)
+                  - 1
+              )
+              / stride[spatial_dim]
+          )
+          + 1
+      )
+
+    output_shape.append(output_dim_size)
+
+  return output_shape
+
+
+def build_transpose_conv(
+    lctx,
+    output_type: ir.RankedTensorType,
+    lhs: ir.Value,
+    rhs: ir.Value,
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    output_padding: list[int],
+    groups: int,
+):
+  lhs_type: ir.RankedTensorType = lhs.type
+  num_spatial_dims = len(lhs_type.shape) - 2
+  rhs = stablehlo.reverse(rhs, list(range(2, 2 + num_spatial_dims)))
+
+  kernel_size = rhs.type.shape
+  # We need to additional padding on the input to get the right output size.
+  adjusted_padding = [
+      dilation[dim] * (kernel_size[dim + 2] - 1) - padding[dim]
+      for dim in range(num_spatial_dims)
+  ]
+  return stablehlo.convolution(
+      result=output_type,
+      lhs=lhs,
+      rhs=rhs,
+      dimension_numbers=create_conv_dimension_numbers(lhs, True),
+      feature_group_count=groups,
+      batch_group_count=1,
+      padding=make_padding(adjusted_padding),
+      lhs_dilation=stride,
+      rhs_dilation=dilation,
+  )
+
+
+# convolution(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride,
+#   SymInt[] padding, SymInt[] dilation, bool transposed,
+#   SymInt[] output_padding, SymInt groups) -> Tensor
+@registry.lower(torch.ops.aten.convolution)
+def _aten_convolution(
+    lctx,
+    lhs: ir.Value,
+    rhs: ir.Value,
+    bias: Optional[ir.Value],
+    stride: list[int],
+    padding: list[int],
+    dilation: list[int],
+    transposed: bool,
+    output_padding: list[int],
+    groups: int,
+):
+
+  # TODO(b/365559296) Add support for output_padding
+  if any(output_padding):
+    raise NotImplementedError(
+        "Output padding on convolution is not implemented."
+    )
+
+  lhs_type: ir.RankedTensorType = lhs.type
+  output_shape = infer_output_shape(
+      lhs, rhs, stride, dilation, padding, transposed, output_padding
+  )
+  output_type = ir.RankedTensorType.get(
+      output_shape,
+      lhs_type.element_type,
+  )
+
+  if transposed:
+    res = build_transpose_conv(
+        lctx,
+        output_type,
+        lhs,
+        rhs,
+        stride,
+        padding,
+        dilation,
+        output_padding,
+        groups,
+    )
+  else:
+    res = stablehlo.convolution(
+        result=output_type,
+        lhs=lhs,
+        rhs=rhs,
+        dimension_numbers=create_conv_dimension_numbers(lhs),
+        feature_group_count=groups,
+        batch_group_count=1,
+        window_strides=stride,
+        padding=make_padding(padding),
+        rhs_dilation=dilation,
+    )
+
+  if bias is not None:
+    # broadcast [C] to [NCHW]
+    broadcasted_bias = stablehlo.broadcast_in_dim(
+        output_type, bias, ir.DenseI64ArrayAttr.get([1])
+    )
+    res = stablehlo.add(
+        lhs=res,
+        rhs=broadcasted_bias,
+    )
+
+  return res

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -0,0 +1,285 @@
+# 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 functools
+import logging
+
+from ai_edge_torch.odml_torch import jax_bridge
+from ai_edge_torch.odml_torch.lowerings import context
+from ai_edge_torch.odml_torch.lowerings import registry
+import jax.numpy as jnp
+from jax._src.lib.mlir import ir
+import torch
+import torch_xla2.ops.jaten  # Import to load torch_xla2 ops
+import torch_xla2.ops.ops_registry  # Import to load torch_xla2 ops
+
+LoweringContext = context.LoweringContext
+
+
+@functools.cache
+def _log_usage(op):
+  logging.warning("Use jax lowering: %s", str(op))
+
+
+def lower_by_jax(op, ir_input_names=None):
+  def inner(lowering):
+    bridged = jax_bridge.wrap(lowering, ir_input_names)
+
+    @registry.lower(op)
+    def _jax_lowering(lctx, *args, **kwargs):
+      _log_usage(op)
+      return bridged(lctx, *args, **kwargs)
+
+    return lowering
+
+  return inner
+
+
+_TORCH_XLA2_IMPLS = {}
+
+for op, torch_xla2_op in torch_xla2.ops.ops_registry.all_aten_ops.items():
+  if not torch_xla2_op.is_jax_function:
+    continue
+  if isinstance(op, torch._ops.OpOverloadPacket):
+    ops = [getattr(op, overload) for overload in op.overloads()] + [op]
+  else:
+    ops = [op]
+
+  for op in ops:
+    _TORCH_XLA2_IMPLS[op] = torch_xla2_op.func
+
+
+def lower_by_torch_xla2(op):
+  return lower_by_jax(op)(_TORCH_XLA2_IMPLS[op])
+
+
+lower_by_torch_xla2(torch.ops.aten._adaptive_avg_pool2d)
+lower_by_torch_xla2(torch.ops.aten._adaptive_avg_pool3d)
+lower_by_torch_xla2(torch.ops.aten._cdist_forward)
+lower_by_torch_xla2(torch.ops.aten._local_scalar_dense)
+lower_by_torch_xla2(torch.ops.aten._local_scalar_dense)
+lower_by_torch_xla2(torch.ops.aten._log_softmax)
+lower_by_torch_xla2(torch.ops.aten._native_batch_norm_legit)
+lower_by_torch_xla2(torch.ops.aten._native_batch_norm_legit_no_training)
+lower_by_torch_xla2(torch.ops.aten._pdist_forward)
+lower_by_torch_xla2(torch.ops.aten._softmax)
+lower_by_torch_xla2(torch.ops.aten._to_copy)
+lower_by_torch_xla2(torch.ops.aten._unsafe_index)
+lower_by_torch_xla2(torch.ops.aten._unsafe_view)
+lower_by_torch_xla2(torch.ops.aten.abs)
+lower_by_torch_xla2(torch.ops.aten.acos)
+lower_by_torch_xla2(torch.ops.aten.acosh)
+lower_by_torch_xla2(torch.ops.aten.add.Scalar)
+lower_by_torch_xla2(torch.ops.aten.add.Tensor)
+lower_by_torch_xla2(torch.ops.aten.addbmm.default)
+lower_by_torch_xla2(torch.ops.aten.addmm)
+lower_by_torch_xla2(torch.ops.aten.addmv)
+lower_by_torch_xla2(torch.ops.aten.alias)
+lower_by_torch_xla2(torch.ops.aten.allclose)
+lower_by_torch_xla2(torch.ops.aten.amax)
+lower_by_torch_xla2(torch.ops.aten.amin)
+lower_by_torch_xla2(torch.ops.aten.any)
+lower_by_torch_xla2(torch.ops.aten.arange.default)
+lower_by_torch_xla2(torch.ops.aten.arange.start)
+lower_by_torch_xla2(torch.ops.aten.arange.start_step)
+lower_by_torch_xla2(torch.ops.aten.argmax)
+lower_by_torch_xla2(torch.ops.aten.argmin)
+lower_by_torch_xla2(torch.ops.aten.as_strided)
+lower_by_torch_xla2(torch.ops.aten.as_strided_copy)
+lower_by_torch_xla2(torch.ops.aten.asin)
+lower_by_torch_xla2(torch.ops.aten.asinh)
+lower_by_torch_xla2(torch.ops.aten.atan)
+lower_by_torch_xla2(torch.ops.aten.atan2)
+lower_by_torch_xla2(torch.ops.aten.atanh)
+lower_by_torch_xla2(torch.ops.aten.avg_pool2d)
+lower_by_torch_xla2(torch.ops.aten.avg_pool3d)
+lower_by_torch_xla2(torch.ops.aten.bitwise_and)
+lower_by_torch_xla2(torch.ops.aten.bitwise_not)
+lower_by_torch_xla2(torch.ops.aten.bitwise_or)
+lower_by_torch_xla2(torch.ops.aten.bitwise_xor)
+lower_by_torch_xla2(torch.ops.aten.bmm)
+lower_by_torch_xla2(torch.ops.aten.ceil)
+lower_by_torch_xla2(torch.ops.aten.clamp.Tensor)
+lower_by_torch_xla2(torch.ops.aten.clamp.default)
+lower_by_torch_xla2(torch.ops.aten.clone)
+lower_by_torch_xla2(torch.ops.aten.clone.default)
+lower_by_torch_xla2(torch.ops.aten.constant_pad_nd)
+lower_by_torch_xla2(torch.ops.aten.cos)
+lower_by_torch_xla2(torch.ops.aten.cosh)
+lower_by_torch_xla2(torch.ops.aten.cumsum)
+lower_by_torch_xla2(torch.ops.aten.detach)
+lower_by_torch_xla2(torch.ops.aten.diagonal)
+lower_by_torch_xla2(torch.ops.aten.div)
+lower_by_torch_xla2(torch.ops.aten.dot)
+lower_by_torch_xla2(torch.ops.aten.embedding)
+lower_by_torch_xla2(torch.ops.aten.empty)
+lower_by_torch_xla2(torch.ops.aten.eq)
+lower_by_torch_xla2(torch.ops.aten.erf)
+lower_by_torch_xla2(torch.ops.aten.exp)
+lower_by_torch_xla2(torch.ops.aten.expand)
+lower_by_torch_xla2(torch.ops.aten.expand_copy)
+lower_by_torch_xla2(torch.ops.aten.expm1)
+lower_by_torch_xla2(torch.ops.aten.fill)
+lower_by_torch_xla2(torch.ops.aten.flip)
+lower_by_torch_xla2(torch.ops.aten.fmod)
+lower_by_torch_xla2(torch.ops.aten.full)
+lower_by_torch_xla2(torch.ops.aten.full_like)
+lower_by_torch_xla2(torch.ops.aten.gather)
+lower_by_torch_xla2(torch.ops.aten.ge)
+lower_by_torch_xla2(torch.ops.aten.gelu)
+lower_by_torch_xla2(torch.ops.aten.glu)
+lower_by_torch_xla2(torch.ops.aten.glu.default)
+lower_by_torch_xla2(torch.ops.aten.gt)
+lower_by_torch_xla2(torch.ops.aten.hardtanh)
+lower_by_torch_xla2(torch.ops.aten.index)
+lower_by_torch_xla2(torch.ops.aten.index.Tensor)
+lower_by_torch_xla2(torch.ops.aten.index_copy)
+lower_by_torch_xla2(torch.ops.aten.index_put)
+lower_by_torch_xla2(torch.ops.aten.index_select)
+lower_by_torch_xla2(torch.ops.aten.isinf)
+lower_by_torch_xla2(torch.ops.aten.isnan)
+lower_by_torch_xla2(torch.ops.aten.le)
+lower_by_torch_xla2(torch.ops.aten.leaky_relu)
+lower_by_torch_xla2(torch.ops.aten.lift_fresh_copy)
+lower_by_torch_xla2(torch.ops.aten.linalg_vector_norm)
+lower_by_torch_xla2(torch.ops.aten.log)
+lower_by_torch_xla2(torch.ops.aten.log10)
+lower_by_torch_xla2(torch.ops.aten.log1p)
+lower_by_torch_xla2(torch.ops.aten.log2)
+lower_by_torch_xla2(torch.ops.aten.logical_and)
+lower_by_torch_xla2(torch.ops.aten.logical_not)
+lower_by_torch_xla2(torch.ops.aten.logical_or)
+lower_by_torch_xla2(torch.ops.aten.logical_xor)
+lower_by_torch_xla2(torch.ops.aten.lt)
+lower_by_torch_xla2(torch.ops.aten.max)
+lower_by_torch_xla2(torch.ops.aten.max_pool2d_with_indices)
+lower_by_torch_xla2(torch.ops.aten.max_pool2d_with_indices_backward)
+lower_by_torch_xla2(torch.ops.aten.max_pool2d_with_indices_backward)
+lower_by_torch_xla2(torch.ops.aten.max_pool3d_with_indices)
+lower_by_torch_xla2(torch.ops.aten.maximum)
+lower_by_torch_xla2(torch.ops.aten.mean)
+lower_by_torch_xla2(torch.ops.aten.min)
+lower_by_torch_xla2(torch.ops.aten.minimum)
+lower_by_torch_xla2(torch.ops.aten.mm)
+lower_by_torch_xla2(torch.ops.aten.mul.Scalar)
+lower_by_torch_xla2(torch.ops.aten.mul.Tensor)
+lower_by_torch_xla2(torch.ops.aten.native_batch_norm)
+lower_by_torch_xla2(torch.ops.aten.native_layer_norm_backward)
+lower_by_torch_xla2(torch.ops.aten.ne)
+lower_by_torch_xla2(torch.ops.aten.neg)
+lower_by_torch_xla2(torch.ops.aten.nonzero)
+lower_by_torch_xla2(torch.ops.aten.outer)
+lower_by_torch_xla2(torch.ops.aten.permute)
+lower_by_torch_xla2(torch.ops.aten.permute_copy)
+lower_by_torch_xla2(torch.ops.aten.pixel_shuffle)
+lower_by_torch_xla2(torch.ops.aten.pow)
+lower_by_torch_xla2(torch.ops.aten.prod)
+lower_by_torch_xla2(torch.ops.aten.reciprocal)
+lower_by_torch_xla2(torch.ops.aten.reflection_pad1d)
+lower_by_torch_xla2(torch.ops.aten.relu)
+lower_by_torch_xla2(torch.ops.aten.remainder)
+lower_by_torch_xla2(torch.ops.aten.repeat)
+lower_by_torch_xla2(torch.ops.aten.reshape)
+lower_by_torch_xla2(torch.ops.aten.roll)
+lower_by_torch_xla2(torch.ops.aten.round)
+lower_by_torch_xla2(torch.ops.aten.rsqrt)
+lower_by_torch_xla2(torch.ops.aten.scalar_tensor)
+lower_by_torch_xla2(torch.ops.aten.scatter.src)
+lower_by_torch_xla2(torch.ops.aten.scatter.value)
+lower_by_torch_xla2(torch.ops.aten.scatter_add)
+lower_by_torch_xla2(torch.ops.aten.scatter_reduce)
+lower_by_torch_xla2(torch.ops.aten.select)
+lower_by_torch_xla2(torch.ops.aten.select_copy)
+lower_by_torch_xla2(torch.ops.aten.select_scatter)
+lower_by_torch_xla2(torch.ops.aten.sigmoid)
+lower_by_torch_xla2(torch.ops.aten.sign)
+lower_by_torch_xla2(torch.ops.aten.silu)
+lower_by_torch_xla2(torch.ops.aten.sin)
+lower_by_torch_xla2(torch.ops.aten.sinh)
+lower_by_torch_xla2(torch.ops.aten.slice)
+lower_by_torch_xla2(torch.ops.aten.slice_copy)
+lower_by_torch_xla2(torch.ops.aten.sort)
+lower_by_torch_xla2(torch.ops.aten.split)
+lower_by_torch_xla2(torch.ops.aten.split_copy)
+lower_by_torch_xla2(torch.ops.aten.split_with_sizes)
+lower_by_torch_xla2(torch.ops.aten.sqrt)
+lower_by_torch_xla2(torch.ops.aten.squeeze)
+lower_by_torch_xla2(torch.ops.aten.squeeze_copy)
+lower_by_torch_xla2(torch.ops.aten.stack)
+lower_by_torch_xla2(torch.ops.aten.sub.Scalar)
+lower_by_torch_xla2(torch.ops.aten.sub.Tensor)
+lower_by_torch_xla2(torch.ops.aten.sum)
+lower_by_torch_xla2(torch.ops.aten.sym_size)
+lower_by_torch_xla2(torch.ops.aten.t)
+lower_by_torch_xla2(torch.ops.aten.tan)
+lower_by_torch_xla2(torch.ops.aten.tanh)
+lower_by_torch_xla2(torch.ops.aten.tensor_split.sections)
+lower_by_torch_xla2(torch.ops.aten.tensor_split.sections)
+lower_by_torch_xla2(torch.ops.aten.to.device)
+lower_by_torch_xla2(torch.ops.aten.to.device)
+lower_by_torch_xla2(torch.ops.aten.to.dtype)
+lower_by_torch_xla2(torch.ops.aten.topk)
+lower_by_torch_xla2(torch.ops.aten.transpose)
+lower_by_torch_xla2(torch.ops.aten.transpose_copy)
+lower_by_torch_xla2(torch.ops.aten.triu)
+lower_by_torch_xla2(torch.ops.aten.true_divide)
+lower_by_torch_xla2(torch.ops.aten.trunc)
+lower_by_torch_xla2(torch.ops.aten.unbind_copy)
+lower_by_torch_xla2(torch.ops.aten.unsqueeze)
+lower_by_torch_xla2(torch.ops.aten.unsqueeze.default)
+lower_by_torch_xla2(torch.ops.aten.unsqueeze_copy)
+lower_by_torch_xla2(torch.ops.aten.var.correction)
+lower_by_torch_xla2(torch.ops.aten.var_mean.correction)
+lower_by_torch_xla2(torch.ops.aten.view)
+lower_by_torch_xla2(torch.ops.aten.view_as_complex)
+lower_by_torch_xla2(torch.ops.aten.view_as_real)
+lower_by_torch_xla2(torch.ops.aten.view_copy)
+lower_by_torch_xla2(torch.ops.aten.where.ScalarOther)
+lower_by_torch_xla2(torch.ops.aten.where.ScalarSelf)
+lower_by_torch_xla2(torch.ops.aten.where.self)
+lower_by_torch_xla2(torch.ops.prims.broadcast_in_dim)
+lower_by_torch_xla2(torch.ops.prims.var)
+
+
+@lower_by_jax(torch.ops.aten.unbind)
+def _aten_copy(self, *args, **kwargs):
+  return _TORCH_XLA2_IMPLS[torch.ops.aten.unbind_copy](self, *args, **kwargs)
+
+
+@lower_by_jax(torch.ops.aten.copy, ir_input_names=["src"])
+def _aten_copy(self, src, **kwargs):
+  return _TORCH_XLA2_IMPLS[torch.ops.aten.copy](self, src)
+
+
+# Schema:
+#   - aten::einsum(str equation, Tensor[] tensors, *, int[]? path=None)
+#       -> Tensor
+# Torch Reference:
+#   - https://pytorch.org/docs/stable/generated/torch.einsum.html
+#   - https://github.com/pytorch/pytorch/blob/1b3f8b75896720e88362cbec7db32abc52afa83e/aten/src/ATen/native/Linear.cpp#L255
+@registry.lower(torch.ops.aten.einsum.default)
+def _aten_einsum_default(
+    lctx: LoweringContext,
+    equation: str,
+    tensors: list[ir.Value],
+    path=None,
+):
+  _log_usage(torch.ops.aten.einsum.default)
+
+  @jax_bridge.wrap
+  def jax_lowering(operands):
+    # Ignore the input path and let JAX determine the path.
+    return jnp.einsum(equation, *operands, optimize="optimal")
+
+  return jax_lowering(lctx, tuple(tensors))

ai_edge_torch/odml_torch/lowerings/_layer_norm.py

@@ -0,0 +1,87 @@
+# 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 stablehlo for LayerNorm."""
+
+import math
+from typing import Optional
+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
+
+
+# native_layer_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight,
+# Tensor? bias, float eps) -> (Tensor, Tensor, Tensor)
+@registry.lower(torch.ops.aten.native_layer_norm)
+def _aten_native_layer_norm(
+    lctx,
+    data: ir.Value,
+    normalized_shape: list[int],
+    weight: Optional[ir.Value],
+    bias: Optional[ir.Value],
+    eps: float,
+):
+  data_type: ir.RankedTensorType = data.type
+  unnormalized_count = math.prod(data_type.shape) // math.prod(normalized_shape)
+  dest_shape = [
+      1,
+      unnormalized_count,
+      math.prod(normalized_shape),
+  ]
+  dest_type = ir.RankedTensorType.get(dest_shape, data_type.element_type)
+
+  reshaped_data = stablehlo.reshape(dest_type, data)
+
+  one = utils.splat(1, data_type.element_type, [unnormalized_count])
+  zero = utils.splat(0, data_type.element_type, [unnormalized_count])
+  output, mean, var = stablehlo.batch_norm_training(
+      reshaped_data, one, zero, eps, 1
+  )
+  eps_splat = utils.splat(eps, var.type.element_type, var.type.shape)
+  rstd = stablehlo.rsqrt(stablehlo.add(var, eps_splat))
+
+  stats_shape = data_type.shape[: -1 * len(normalized_shape)] + [1] * len(
+      normalized_shape
+  )
+  stats_type = ir.RankedTensorType.get(stats_shape, data_type.element_type)
+  mean = stablehlo.reshape(stats_type, mean)
+  rstd = stablehlo.reshape(stats_type, rstd)
+
+  output = stablehlo.reshape(data_type, output)
+
+  data_rank = len(data_type.shape)
+  normalized_rank = len(normalized_shape)
+  if weight is not None:
+    weight = stablehlo.broadcast_in_dim(
+        data_type,
+        weight,
+        ir.DenseI64ArrayAttr.get(
+            list(range(data_rank - normalized_rank, data_rank))
+        ),
+    )
+    output = stablehlo.multiply(weight, output)
+  if bias is not None:
+    bias = stablehlo.broadcast_in_dim(
+        data_type,
+        bias,
+        ir.DenseI64ArrayAttr.get(
+            list(range(data_rank - normalized_rank, data_rank))
+        ),
+    )
+    output = stablehlo.add(bias, output)
+
+  return output, mean, rstd