
# Autogenerated by mlir-tblgen; don't manually edit.

from ._ods_common import _cext as _ods_cext
from ._ods_common import (
    equally_sized_accessor as _ods_equally_sized_accessor,
    get_default_loc_context as _ods_get_default_loc_context,
    get_op_results_or_values as _get_op_results_or_values,
    segmented_accessor as _ods_segmented_accessor,
)
_ods_ir = _ods_cext.ir
_ods_cext.globals.register_traceback_file_exclusion(__file__)

import builtins
from typing import Any as _Any, Sequence as _Sequence, Union as _Union, Optional as _Optional
import sys as _sys
if _sys.version_info >= (3, 12):
  from collections.abc import Buffer as _Buffer  # pytype: disable=not-supported-yet
else:
  try:
    from typing_extensions import Buffer as _Buffer
  except ImportError:
    _Buffer = _Any


@_ods_cext.register_dialect
class _Dialect(_ods_ir.Dialect):
  DIALECT_NAMESPACE = "shard"

@_ods_cext.register_operation(_Dialect)
class AllGatherOp(_ods_ir.OpView):
  r"""
  Concatenates all tensor slices from a device group defined by `grid_axes` along
  the tensor dimension `gather_axis` and replicates the result across all devices
  in the group.
  
  Example:
  ```mlir
  shard.grid @grid0(shape = 2x2)
  ...
  %1 = shard.all_gather %0 on @grid0 grid_axes = [1] gather_axis = 1
    : tensor<2x2xi8> -> tensor<2x4xi8>
  ```
  Input:
  ```
                   +-------+-------+
  device (0, 0) -> |  1  2 |  5  6 | <- device (0, 1)
                   |  3  4 |  7  8 |
                   +-------+-------+
  device (1, 0) -> |  9 10 | 13 14 | <- device (1, 1)
                   | 11 12 | 15 16 |
                   +-------+-------+
  ```
  Result:
  ```
  gather tensor
  axis 1
  ------------>
  +-------------+
  |  1  2  5  6 | <- devices (0, 0) and (0, 1)
  |  3  4  7  8 |
  +-------------+
  |  9 10 13 14 | <- devices (1, 0) and (1, 1)
  | 11 12 15 16 |
  +-------------+
  ```
  """

  OPERATION_NAME = "shard.all_gather"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, gather_axis: _Union[int, _ods_ir.IntegerAttr], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(input)
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    if grid_axes is not None: attributes["grid_axes"] = (grid_axes if (
        isinstance(grid_axes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
          _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(grid_axes, context=_ods_context))
    attributes["gather_axis"] = (gather_axis if (
    isinstance(gather_axis, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('IndexAttr')) else
      _ods_ir.AttrBuilder.get('IndexAttr')(gather_axis, context=_ods_context))
    results = []
    results.append(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operation.operands[0]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["grid_axes"]

  @grid_axes.setter
  def grid_axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid_axes"] = value

  @builtins.property
  def gather_axis(self) -> _ods_ir.IntegerAttr:
    return self.operation.attributes["gather_axis"]

  @gather_axis.setter
  def gather_axis(self, value: _ods_ir.IntegerAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["gather_axis"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResult:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(AllGatherOp)
class AllGatherOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.all_gather"

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operands[0]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["grid_axes"]

  @builtins.property
  def gather_axis(self) -> _ods_ir.IntegerAttr:
    return self.attributes["gather_axis"]

def all_gather(result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, gather_axis: _Union[int, _ods_ir.IntegerAttr], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult:
  return AllGatherOp(result=result, grid=grid, input=input, gather_axis=gather_axis, grid_axes=grid_axes, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class AllReduceOp(_ods_ir.OpView):
  r"""
  Reduces the input tensor across all devices within the groups defined by
  `grid_axes`, using the specified reduction method. The operation performs an
  element-wise reduction over the tensor slices from all devices in each group.
  Each device in a group receives a replicated copy of the reduction result.
  The accumulation element type is determined by the result type and does not
  need to match the input element type. Before performing the reduction, each
  input element is converted to the result element type.
  
  Attributes:
  `reduction`: Indicates the reduction method.
  
  Example:
  ```
  %1 = shard.all_reduce %0 on @grid0 grid_axes = [1, 0] reduction = <max>
    : tensor<3x4xf32> -> tensor<3x4xf64>
  ```
  """

  OPERATION_NAME = "shard.all_reduce"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, reduction: _Optional[_Union[_Any, _ods_ir.Attribute]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(input)
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    if grid_axes is not None: attributes["grid_axes"] = (grid_axes if (
        isinstance(grid_axes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
          _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(grid_axes, context=_ods_context))
    if reduction is not None: attributes["reduction"] = (reduction if (
        isinstance(reduction, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('shard.Shard_ReductionKindAttr')) else
          _ods_ir.AttrBuilder.get('shard.Shard_ReductionKindAttr')(reduction, context=_ods_context))
    results = []
    results.append(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operation.operands[0]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["grid_axes"]

  @grid_axes.setter
  def grid_axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid_axes"] = value

  @builtins.property
  def reduction(self) -> _ods_ir.Attribute:
    return self.operation.attributes["reduction"]

  @reduction.setter
  def reduction(self, value: _ods_ir.Attribute):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["reduction"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResult:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(AllReduceOp)
class AllReduceOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.all_reduce"

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operands[0]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["grid_axes"]

  @builtins.property
  def reduction(self) -> _ods_ir.Attribute:
    return self.attributes["reduction"]

def all_reduce(result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, reduction: _Optional[_Union[_Any, _ods_ir.Attribute]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult:
  return AllReduceOp(result=result, grid=grid, input=input, grid_axes=grid_axes, reduction=reduction, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class AllSliceOp(_ods_ir.OpView):
  r"""
  Within each device group defined by `grid_axes`, slices the input tensor along
  the `slice_axis` dimension. It can be viewed as the inverse of an all-gather if
  the input data is replicated along the `slice_axis`.
  Each process simply crops its local data to the slice corresponding to its
  in-group device index.
  Notice: `AllSliceOp` does not involve any communication between devices and
          devices within a group may not have replicated input data.
  
  Example:
  ```mlir
  shard.grid @grid0(shape = 2x2)
  ...
  %1 = shard.all_slice %0 on @grid0 grid_axes = [1] slice_axis = 1
    : tensor<2x4xi8> -> tensor<2x2xi8>
  ```
  Input:
  ```
  +-------------+
  |  1  2  5  6 | <- devices (0, 0) and (0, 1)
  |  3  4  7  8 |
  +-------------+
  |  9 10 13 14 | <- devices (1, 0) and (1, 1)
  | 11 12 15 16 |
  +-------------+
  ```
  Result:
  ```
  slice tensor
  axis 1
  ------------>
                   +-------+-------+
  device (0, 0) -> |  1  2 |  5  6 | <- device (0, 1)
                   |  3  4 |  7  8 |
                   +-------+-------+
  device (1, 0) -> |  9 10 | 13 14 | <- device (1, 1)
                   | 11 12 | 15 16 |
                   +-------+-------+
  ```
  """

  OPERATION_NAME = "shard.all_slice"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, slice_axis: _Union[int, _ods_ir.IntegerAttr], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(input)
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    if grid_axes is not None: attributes["grid_axes"] = (grid_axes if (
        isinstance(grid_axes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
          _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(grid_axes, context=_ods_context))
    attributes["slice_axis"] = (slice_axis if (
    isinstance(slice_axis, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('IndexAttr')) else
      _ods_ir.AttrBuilder.get('IndexAttr')(slice_axis, context=_ods_context))
    results = []
    results.append(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operation.operands[0]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["grid_axes"]

  @grid_axes.setter
  def grid_axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid_axes"] = value

  @builtins.property
  def slice_axis(self) -> _ods_ir.IntegerAttr:
    return self.operation.attributes["slice_axis"]

  @slice_axis.setter
  def slice_axis(self, value: _ods_ir.IntegerAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["slice_axis"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResult:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(AllSliceOp)
class AllSliceOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.all_slice"

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operands[0]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["grid_axes"]

  @builtins.property
  def slice_axis(self) -> _ods_ir.IntegerAttr:
    return self.attributes["slice_axis"]

def all_slice(result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, slice_axis: _Union[int, _ods_ir.IntegerAttr], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult:
  return AllSliceOp(result=result, grid=grid, input=input, slice_axis=slice_axis, grid_axes=grid_axes, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class AllToAllOp(_ods_ir.OpView):
  r"""
  Each participant logically splits its input along split_axis,
  then scatters the resulting pieces across the group defined by `grid_axes`.
  After receiving data pieces from other participants' scatters,
  it concatenates them along concat_axis to produce the final result.
  
  Example:
  ```
  shard.grid @grid0(shape = 3)
  ...
  %1 = shard.all_to_all %0 on @grid0 grid_axes = [0]
    split_axis = 0 concat_axis = 0
    : tensor<3x2xi8> -> tensor<3x2xi8>
  ```
  Input:
  ```
   device  device  device
   (0)     (1)     (2)
  +-------+-------+-------+  | split and concat along
  | 11 12 | 21 22 | 31 32 |  | tensor axis 0
  | 13 14 | 23 24 | 33 34 |  ↓
  | 15 16 | 25 26 | 35 36 |
  +-------+-------+-------+
  ```
  Result:
  ```
   device  device  device
   (0)     (1)     (2)
  +-------+-------+-------+
  | 11 12 | 13 14 | 15 16 |
  | 21 22 | 23 24 | 25 26 |
  | 31 32 | 33 34 | 35 36 |
  +-------+-------+-------+
  ```
  """

  OPERATION_NAME = "shard.all_to_all"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, split_axis: _Union[int, _ods_ir.IntegerAttr], concat_axis: _Union[int, _ods_ir.IntegerAttr], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(input)
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    if grid_axes is not None: attributes["grid_axes"] = (grid_axes if (
        isinstance(grid_axes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
          _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(grid_axes, context=_ods_context))
    attributes["split_axis"] = (split_axis if (
    isinstance(split_axis, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('IndexAttr')) else
      _ods_ir.AttrBuilder.get('IndexAttr')(split_axis, context=_ods_context))
    attributes["concat_axis"] = (concat_axis if (
    isinstance(concat_axis, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('IndexAttr')) else
      _ods_ir.AttrBuilder.get('IndexAttr')(concat_axis, context=_ods_context))
    results = []
    results.append(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operation.operands[0]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["grid_axes"]

  @grid_axes.setter
  def grid_axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid_axes"] = value

  @builtins.property
  def split_axis(self) -> _ods_ir.IntegerAttr:
    return self.operation.attributes["split_axis"]

  @split_axis.setter
  def split_axis(self, value: _ods_ir.IntegerAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["split_axis"] = value

  @builtins.property
  def concat_axis(self) -> _ods_ir.IntegerAttr:
    return self.operation.attributes["concat_axis"]

  @concat_axis.setter
  def concat_axis(self, value: _ods_ir.IntegerAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["concat_axis"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResult:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(AllToAllOp)
class AllToAllOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.all_to_all"

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operands[0]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["grid_axes"]

  @builtins.property
  def split_axis(self) -> _ods_ir.IntegerAttr:
    return self.attributes["split_axis"]

  @builtins.property
  def concat_axis(self) -> _ods_ir.IntegerAttr:
    return self.attributes["concat_axis"]

def all_to_all(result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, split_axis: _Union[int, _ods_ir.IntegerAttr], concat_axis: _Union[int, _ods_ir.IntegerAttr], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult:
  return AllToAllOp(result=result, grid=grid, input=input, split_axis=split_axis, concat_axis=concat_axis, grid_axes=grid_axes, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class BroadcastOp(_ods_ir.OpView):
  r"""
  Copies the input tensor on `root` to all devices in each group defined by
  `grid_axes`. The `root` device is defined by its in-group multi-index.
  The contents of input tensors on non-root devices are ignored.
  
  Example:
  ```
  shard.grid @grid0(shape = 2x2)
  
  %1 = shard.broadcast %0 on @grid0
    grid_axes = [0]
    root = [0]
    : (tensor<2xi8>) -> tensor<2xi8>
  ```
  
  Input:
  ```
                   +-------+-------+                   | broadcast
  device (0, 0) -> |  1  2 |  3  4 | <- device (0, 1)  | along axis 0
                   +-------+-------+                   ↓
  device (1, 0) -> |  *  * |  *  * | <- device (1, 1)
                   +-------+-------+
  ```
  
  Output:
  ```
                   +-------+-------+
  device (0, 0) -> |  1  2 |  3  4 | <- device (0, 1)
                   +-------+-------+
  device (1, 0) -> |  1  2 |  3  4 | <- device (1, 1)
                   +-------+-------+
  ```
  """

  OPERATION_NAME = "shard.broadcast"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value[_ods_ir.RankedTensorType], root: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], root_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(input)
    operands.extend(_get_op_results_or_values(root_dynamic))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    if grid_axes is not None: attributes["grid_axes"] = (grid_axes if (
        isinstance(grid_axes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
          _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(grid_axes, context=_ods_context))
    attributes["root"] = (root if (
    isinstance(root, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('DenseI64ArrayAttr')) else
      _ods_ir.AttrBuilder.get('DenseI64ArrayAttr')(root, context=_ods_context))
    results = []
    results.append(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def input(self) -> _ods_ir.Value[_ods_ir.RankedTensorType]:
    return self.operation.operands[0]

  @builtins.property
  def root_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operation.operands) - 2 + 1
    return self.operation.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["grid_axes"]

  @grid_axes.setter
  def grid_axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid_axes"] = value

  @builtins.property
  def root(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.operation.attributes["root"]

  @root.setter
  def root(self, value: _ods_ir.DenseI64ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["root"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResult[_ods_ir.RankedTensorType]:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(BroadcastOp)
class BroadcastOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.broadcast"

  @builtins.property
  def input(self) -> _ods_ir.Value[_ods_ir.RankedTensorType]:
    return self.operands[0]

  @builtins.property
  def root_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operands) - 2 + 1
    return self.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["grid_axes"]

  @builtins.property
  def root(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.attributes["root"]

def broadcast(result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value[_ods_ir.RankedTensorType], root: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], root_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult[_ods_ir.RankedTensorType]:
  return BroadcastOp(result=result, grid=grid, input=input, root=root, root_dynamic=root_dynamic, grid_axes=grid_axes, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class GatherOp(_ods_ir.OpView):
  r"""
  Concatenates all tensor slices from a device group defined by `grid_axes` along
  the tensor dimension `gather_axis` and returns the resulting tensor on each
  `root` device. The result on all other (non-root) devices is undefined.
  The `root` device is defined by its in-group multi-index.
  
  Example:
  ```mlir
  shard.grid @grid0(shape = 2x2)
  ...
  %1 = shard.gather %0 on @grid0 grid_axes = [1]
    gather_axis = 1 root = [1]
    : (tensor<2x2xi8>) -> tensor<2x4xi8>
  ```
  Input:
  ```
                    gather tensor
                    axis 1
                    ------------>
                   +-------+-------+
  device (0, 0) -> |  1  2 |  5  6 | <- device (0, 1)
                   |  3  4 |  7  8 |
                   +-------+-------+
  device (1, 0) -> |  9 10 | 13 14 | <- device (1, 1)
                   | 11 12 | 15 16 |
                   +-------+-------+
  ```
  Result:
  ```
  +-------------+
  |  1  2  5  6 | <- devices (0, 1)
  |  3  4  7  8 |
  +-------------+
  |  9 10 13 14 | <- devices (1, 1)
  | 11 12 15 16 |
  +-------------+
  ```
  Devices `(0, 0)` and `(1, 0)` have undefined result.
  """

  OPERATION_NAME = "shard.gather"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, gather_axis: _Union[int, _ods_ir.IntegerAttr], root: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], root_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(input)
    operands.extend(_get_op_results_or_values(root_dynamic))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    if grid_axes is not None: attributes["grid_axes"] = (grid_axes if (
        isinstance(grid_axes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
          _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(grid_axes, context=_ods_context))
    attributes["gather_axis"] = (gather_axis if (
    isinstance(gather_axis, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('IndexAttr')) else
      _ods_ir.AttrBuilder.get('IndexAttr')(gather_axis, context=_ods_context))
    attributes["root"] = (root if (
    isinstance(root, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('DenseI64ArrayAttr')) else
      _ods_ir.AttrBuilder.get('DenseI64ArrayAttr')(root, context=_ods_context))
    results = []
    results.append(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operation.operands[0]

  @builtins.property
  def root_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operation.operands) - 2 + 1
    return self.operation.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["grid_axes"]

  @grid_axes.setter
  def grid_axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid_axes"] = value

  @builtins.property
  def gather_axis(self) -> _ods_ir.IntegerAttr:
    return self.operation.attributes["gather_axis"]

  @gather_axis.setter
  def gather_axis(self, value: _ods_ir.IntegerAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["gather_axis"] = value

  @builtins.property
  def root(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.operation.attributes["root"]

  @root.setter
  def root(self, value: _ods_ir.DenseI64ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["root"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResult:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(GatherOp)
class GatherOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.gather"

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operands[0]

  @builtins.property
  def root_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operands) - 2 + 1
    return self.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["grid_axes"]

  @builtins.property
  def gather_axis(self) -> _ods_ir.IntegerAttr:
    return self.attributes["gather_axis"]

  @builtins.property
  def root(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.attributes["root"]

def gather(result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, gather_axis: _Union[int, _ods_ir.IntegerAttr], root: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], root_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult:
  return GatherOp(result=result, grid=grid, input=input, gather_axis=gather_axis, root=root, root_dynamic=root_dynamic, grid_axes=grid_axes, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class GetShardingOp(_ods_ir.OpView):
  r"""
  This operation returns the sharding of the given tensor as a Sharding.
  """

  OPERATION_NAME = "shard.get_sharding"

  _ODS_REGIONS = (0, True)

  def __init__(self, source: _ods_ir.Value[_ods_ir.RankedTensorType], *, results: _Optional[_Sequence[_ods_ir.Type]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(source)
    _ods_context = _ods_get_default_loc_context(loc)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def source(self) -> _ods_ir.Value[_ods_ir.RankedTensorType]:
    return self.operation.operands[0]

  @builtins.property
  def result(self) -> _ods_ir.OpResult:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(GetShardingOp)
class GetShardingOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.get_sharding"

  @builtins.property
  def source(self) -> _ods_ir.Value[_ods_ir.RankedTensorType]:
    return self.operands[0]

def get_sharding(source: _ods_ir.Value[_ods_ir.RankedTensorType], *, results: _Optional[_Sequence[_ods_ir.Type]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult:
  return GetShardingOp(source=source, results=results, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class GridOp(_ods_ir.OpView):
  r"""
  The shard.grid operation is a symbol operation that identifies a specific
  grid. The operation has three attributes:
  
  1. `sym_name`: This attribute uniquely identifies the name of the grid.
  This name serves as a symbolic reference to the grid throughout
  the MLIR module, allowing for consistent referencing and easier debugging.
  
  2. `shape`: This attribute represents the shape of the device grid.
  It uses the same notation as a tensor shape. Also allowing for dynamic
  dimensions.
  This flexibility allows for dynamic device assignment or configurations
  where the exact number of devices might not be determined during compile
  time.
  For example `2x?x4`.
  
  Example:
  ```
  // A device grid with 3 axes, the total device number is 4 * 8 * 12
  // The dimension sizes are 4, 8, 12 
  shard.grid @grid0(shape = 4x8x12)
  
  // A device grid with 2 axes, the total device number is unknown
  // The first dimension size is 4 and the second is unknown
  shard.grid @grid1(shape = 4x?)
  
  // A device grid with 2 axes, the total device number is unknown
  // The first dimension size is unknown and the second is 4
  shard.grid @grid2(shape = ?x4)
  
  // A device grid with 2 axes, the number of devices along both axes
  // is unknown
  shard.grid @grid3(shape = ?x?)
  ```
  """

  OPERATION_NAME = "shard.grid"

  _ODS_REGIONS = (0, True)

  def __init__(self, sym_name: _Union[str, _ods_ir.StringAttr], shape: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], *, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["sym_name"] = (sym_name if (
    isinstance(sym_name, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('SymbolNameAttr')) else
      _ods_ir.AttrBuilder.get('SymbolNameAttr')(sym_name, context=_ods_context))
    attributes["shape"] = (shape if (
    isinstance(shape, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('DenseI64ArrayAttr')) else
      _ods_ir.AttrBuilder.get('DenseI64ArrayAttr')(shape, context=_ods_context))
    results = []
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def sym_name(self) -> _ods_ir.StringAttr:
    return self.operation.attributes["sym_name"]

  @sym_name.setter
  def sym_name(self, value: _ods_ir.StringAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["sym_name"] = value

  @builtins.property
  def shape(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.operation.attributes["shape"]

  @shape.setter
  def shape(self, value: _ods_ir.DenseI64ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["shape"] = value

@_ods_cext.register_op_adaptor(GridOp)
class GridOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.grid"

  @builtins.property
  def sym_name(self) -> _ods_ir.StringAttr:
    return self.attributes["sym_name"]

  @builtins.property
  def shape(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.attributes["shape"]

def grid(sym_name: _Union[str, _ods_ir.StringAttr], shape: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], *, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> GridOp:
  return GridOp(sym_name=sym_name, shape=shape, loc=loc, ip=ip)

@_ods_cext.register_operation(_Dialect)
class GridShapeOp(_ods_ir.OpView):
  OPERATION_NAME = "shard.grid_shape"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _Sequence[_ods_ir.Type], grid: _Union[str, _ods_ir.FlatSymbolRefAttr], *, axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    if axes is not None: attributes["axes"] = (axes if (
        isinstance(axes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
          _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(axes, context=_ods_context))
    results = []
    results.extend(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["axes"]

  @axes.setter
  def axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["axes"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResultList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operation.results) - 1 + 1
    return self.operation.results[0:0 + _ods_variadic_group_length]

@_ods_cext.register_op_adaptor(GridShapeOp)
class GridShapeOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.grid_shape"

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["axes"]

def grid_shape(result: _Sequence[_ods_ir.Type], grid: _Union[str, _ods_ir.FlatSymbolRefAttr], *, axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _Union[_ods_ir.OpResult, _ods_ir.OpResultList, GridShapeOp]:
  op = GridShapeOp(result=result, grid=grid, axes=axes, loc=loc, ip=ip); results = op.results
  return results if len(results) > 1 else (results[0] if len(results) == 1 else op)

@_ods_cext.register_operation(_Dialect)
class NeighborsLinearIndicesOp(_ods_ir.OpView):
  r"""
  Example:
  ```
  shard.grid @grid0(shape = 10x20x30)
  %c1 = arith.constant 1 : index
  %c2 = arith.constant 2 : index
  %c3 = arith.constant 3 : index
  %idx = shard.neighbors_linear_indices on @grid[%c1, %c2, %c3] split_axes = [1] : index
  ```
  The above returns two indices, `633` and `693`, which correspond to the
  index of the previous process `(1, 1, 3)`, and the next process 
  `(1, 3, 3)` along the split axis `1`.
  
  A negative value is returned if there is no neighbor in the respective
  direction along the given `split_axes`.
  """

  OPERATION_NAME = "shard.neighbors_linear_indices"

  _ODS_REGIONS = (0, True)

  def __init__(self, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], device: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], split_axes: _Union[_Any, _ods_ir.DenseI16ArrayAttr], *, results: _Optional[_Sequence[_ods_ir.Type]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.extend(_get_op_results_or_values(device))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    attributes["split_axes"] = (split_axes if (
    isinstance(split_axes, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
      _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(split_axes, context=_ods_context))
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def device(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operation.operands) - 1 + 1
    return self.operation.operands[0:0 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def split_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["split_axes"]

  @split_axes.setter
  def split_axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["split_axes"] = value

  @builtins.property
  def neighbor_down(self) -> _ods_ir.OpResult[_ods_ir.IndexType]:
    return self.operation.results[0]

  @builtins.property
  def neighbor_up(self) -> _ods_ir.OpResult[_ods_ir.IndexType]:
    return self.operation.results[1]

@_ods_cext.register_op_adaptor(NeighborsLinearIndicesOp)
class NeighborsLinearIndicesOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.neighbors_linear_indices"

  @builtins.property
  def device(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operands) - 1 + 1
    return self.operands[0:0 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def split_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["split_axes"]

def neighbors_linear_indices(grid: _Union[str, _ods_ir.FlatSymbolRefAttr], device: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], split_axes: _Union[_Any, _ods_ir.DenseI16ArrayAttr], *, results: _Optional[_Sequence[_ods_ir.Type]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResultList:
  return NeighborsLinearIndicesOp(grid=grid, device=device, split_axes=split_axes, results=results, loc=loc, ip=ip).results

@_ods_cext.register_operation(_Dialect)
class ProcessLinearIndexOp(_ods_ir.OpView):
  r"""
  Example:
  ```
  %idx = shard.process_linear_index on @grid : index
  ```
  if `@grid` has shape `(10, 20, 30)`, a device with multi
  index `(1, 2, 3)` will have linear index `3 + 30*2 + 20*30*1`.
  """

  OPERATION_NAME = "shard.process_linear_index"

  _ODS_REGIONS = (0, True)

  def __init__(self, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], *, results: _Optional[_Sequence[_ods_ir.Type]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResult[_ods_ir.IndexType]:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(ProcessLinearIndexOp)
class ProcessLinearIndexOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.process_linear_index"

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

def process_linear_index(grid: _Union[str, _ods_ir.FlatSymbolRefAttr], *, results: _Optional[_Sequence[_ods_ir.Type]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult[_ods_ir.IndexType]:
  return ProcessLinearIndexOp(grid=grid, results=results, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class ProcessMultiIndexOp(_ods_ir.OpView):
  r"""
  It is used in the SPMD format of IR.
  The `axes` mush be non-negative and less than the total number of grid axes.
  If the axes are empty then get the index along all axes.
  """

  OPERATION_NAME = "shard.process_multi_index"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _Sequence[_ods_ir.Type], grid: _Union[str, _ods_ir.FlatSymbolRefAttr], *, axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    if axes is not None: attributes["axes"] = (axes if (
        isinstance(axes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
          _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(axes, context=_ods_context))
    results = []
    results.extend(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["axes"]

  @axes.setter
  def axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["axes"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResultList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operation.results) - 1 + 1
    return self.operation.results[0:0 + _ods_variadic_group_length]

@_ods_cext.register_op_adaptor(ProcessMultiIndexOp)
class ProcessMultiIndexOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.process_multi_index"

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["axes"]

def process_multi_index(result: _Sequence[_ods_ir.Type], grid: _Union[str, _ods_ir.FlatSymbolRefAttr], *, axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _Union[_ods_ir.OpResult, _ods_ir.OpResultList, ProcessMultiIndexOp]:
  op = ProcessMultiIndexOp(result=result, grid=grid, axes=axes, loc=loc, ip=ip); results = op.results
  return results if len(results) > 1 else (results[0] if len(results) == 1 else op)

@_ods_cext.register_operation(_Dialect)
class RecvOp(_ods_ir.OpView):
  r"""
  Receive tensor from device `source`, which is defined by its in-group
  multi-index. The groups are defined by `grid_axes`.
  The content of input tensor is ignored.
  """

  OPERATION_NAME = "shard.recv"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, source_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, source: _Optional[_Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(input)
    operands.extend(_get_op_results_or_values(source_dynamic))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    if grid_axes is not None: attributes["grid_axes"] = (grid_axes if (
        isinstance(grid_axes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
          _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(grid_axes, context=_ods_context))
    if source is not None: attributes["source"] = (source if (
        isinstance(source, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('DenseI64ArrayAttr')) else
          _ods_ir.AttrBuilder.get('DenseI64ArrayAttr')(source, context=_ods_context))
    results = []
    results.append(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operation.operands[0]

  @builtins.property
  def source_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operation.operands) - 2 + 1
    return self.operation.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["grid_axes"]

  @grid_axes.setter
  def grid_axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid_axes"] = value

  @builtins.property
  def source(self) -> _Optional[_ods_ir.DenseI64ArrayAttr]:
    if "source" not in self.operation.attributes:
      return None
    return self.operation.attributes["source"]

  @source.setter
  def source(self, value: _Optional[_ods_ir.DenseI64ArrayAttr]):
    if value is not None:
      self.operation.attributes["source"] = value
    elif "source" in self.operation.attributes:
      del self.operation.attributes["source"]

  @source.deleter
  def source(self):
    del self.operation.attributes["source"]

  @builtins.property
  def result(self) -> _ods_ir.OpResult[_ods_ir.RankedTensorType]:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(RecvOp)
class RecvOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.recv"

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operands[0]

  @builtins.property
  def source_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operands) - 2 + 1
    return self.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["grid_axes"]

  @builtins.property
  def source(self) -> _Optional[_ods_ir.DenseI64ArrayAttr]:
    if "source" not in self.attributes:
      return None
    return self.attributes["source"]

def recv(result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, source_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, source: _Optional[_Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult[_ods_ir.RankedTensorType]:
  return RecvOp(result=result, grid=grid, input=input, source_dynamic=source_dynamic, grid_axes=grid_axes, source=source, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class ReduceOp(_ods_ir.OpView):
  r"""
  Reduces the input tensor across all devices within the groups defined by
  `grid_axes`, using the specified reduction method. The operation performs an
  element-wise reduction over the tensor slices from all devices in each group.
  The reduction result will be returned on the `root` device of each group.
  It is undefined on all other (non-root) devices.
  The `root` device is defined by its in-group multi-index.
  The accumulation element type is determined by the result type and does not
  need to match the input element type. Before performing the reduction, each
  input element is converted to the result element type.
  
  Attributes:
  `reduction`: Indicates the reduction method.
  
  Example:
  ```
  %1 = shard.reduce %0 on @grid0 grid_axes = [1, 0]
    reduction = <max> root = [2, 3]
    : (tensor<3x4xf32>) -> tensor<3x4xf64>
  ```
  """

  OPERATION_NAME = "shard.reduce"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value[_ods_ir.RankedTensorType], root: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], root_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, reduction: _Optional[_Union[_Any, _ods_ir.Attribute]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(input)
    operands.extend(_get_op_results_or_values(root_dynamic))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    if grid_axes is not None: attributes["grid_axes"] = (grid_axes if (
        isinstance(grid_axes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
          _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(grid_axes, context=_ods_context))
    if reduction is not None: attributes["reduction"] = (reduction if (
        isinstance(reduction, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('shard.Shard_ReductionKindAttr')) else
          _ods_ir.AttrBuilder.get('shard.Shard_ReductionKindAttr')(reduction, context=_ods_context))
    attributes["root"] = (root if (
    isinstance(root, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('DenseI64ArrayAttr')) else
      _ods_ir.AttrBuilder.get('DenseI64ArrayAttr')(root, context=_ods_context))
    results = []
    results.append(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def input(self) -> _ods_ir.Value[_ods_ir.RankedTensorType]:
    return self.operation.operands[0]

  @builtins.property
  def root_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operation.operands) - 2 + 1
    return self.operation.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["grid_axes"]

  @grid_axes.setter
  def grid_axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid_axes"] = value

  @builtins.property
  def reduction(self) -> _ods_ir.Attribute:
    return self.operation.attributes["reduction"]

  @reduction.setter
  def reduction(self, value: _ods_ir.Attribute):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["reduction"] = value

  @builtins.property
  def root(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.operation.attributes["root"]

  @root.setter
  def root(self, value: _ods_ir.DenseI64ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["root"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResult[_ods_ir.RankedTensorType]:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(ReduceOp)
class ReduceOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.reduce"

  @builtins.property
  def input(self) -> _ods_ir.Value[_ods_ir.RankedTensorType]:
    return self.operands[0]

  @builtins.property
  def root_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operands) - 2 + 1
    return self.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["grid_axes"]

  @builtins.property
  def reduction(self) -> _ods_ir.Attribute:
    return self.attributes["reduction"]

  @builtins.property
  def root(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.attributes["root"]

def reduce(result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value[_ods_ir.RankedTensorType], root: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], root_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, reduction: _Optional[_Union[_Any, _ods_ir.Attribute]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult[_ods_ir.RankedTensorType]:
  return ReduceOp(result=result, grid=grid, input=input, root=root, root_dynamic=root_dynamic, grid_axes=grid_axes, reduction=reduction, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class ReduceScatterOp(_ods_ir.OpView):
  r"""
  Reduces the input tensor across all devices within the groups defined by
  `grid_axes` using the specified reduction method. The reduction is performed
  element-wise across the tensor pieces from all devices in the group.
  After reduction, the reduction result is scattered (split and distributed)
  across the device group along `scatter_dim`.
  Example:
  ```
  shard.grid @grid0(shape = 2x2)
  ...
  %1 = shard.reduce_scatter %0 on @grid0 grid_axes = [1]
    reduction = <max> scatter_dim = 0
    : tensor<2x2xf32> -> tensor<1x2xf64>
  ```
  Input:
  ```
                            device
                            (0, 1)
                               ↓
                   +-------+-------+  | scatter tensor
  device (0, 0) -> |  1  2 |  5  6 |  | axis 0
                   |  3  4 |  7  8 |  ↓
                   +-------+-------+
  device (1, 0) -> |  9 10 | 13 14 |
                   | 11 12 | 15 16 |
                   +-------+-------+
                              ↑
                            device
                            (1, 1)
  ```
  Result:
  ```
  +-------+
  |  5  6 | <- devices (0, 0)
  +-------+
  |  7  8 | <- devices (0, 1)
  +-------+
  | 13 14 | <- devices (1, 0)
  +-------+
  | 15 16 | <- devices (1, 1)
  +-------+
  ```
  """

  OPERATION_NAME = "shard.reduce_scatter"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, scatter_dim: _Union[int, _ods_ir.IntegerAttr], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, reduction: _Optional[_Union[_Any, _ods_ir.Attribute]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(input)
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    if grid_axes is not None: attributes["grid_axes"] = (grid_axes if (
        isinstance(grid_axes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
          _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(grid_axes, context=_ods_context))
    if reduction is not None: attributes["reduction"] = (reduction if (
        isinstance(reduction, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('shard.Shard_ReductionKindAttr')) else
          _ods_ir.AttrBuilder.get('shard.Shard_ReductionKindAttr')(reduction, context=_ods_context))
    attributes["scatter_dim"] = (scatter_dim if (
    isinstance(scatter_dim, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('IndexAttr')) else
      _ods_ir.AttrBuilder.get('IndexAttr')(scatter_dim, context=_ods_context))
    results = []
    results.append(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operation.operands[0]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["grid_axes"]

  @grid_axes.setter
  def grid_axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid_axes"] = value

  @builtins.property
  def reduction(self) -> _ods_ir.Attribute:
    return self.operation.attributes["reduction"]

  @reduction.setter
  def reduction(self, value: _ods_ir.Attribute):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["reduction"] = value

  @builtins.property
  def scatter_dim(self) -> _ods_ir.IntegerAttr:
    return self.operation.attributes["scatter_dim"]

  @scatter_dim.setter
  def scatter_dim(self, value: _ods_ir.IntegerAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["scatter_dim"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResult:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(ReduceScatterOp)
class ReduceScatterOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.reduce_scatter"

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operands[0]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["grid_axes"]

  @builtins.property
  def reduction(self) -> _ods_ir.Attribute:
    return self.attributes["reduction"]

  @builtins.property
  def scatter_dim(self) -> _ods_ir.IntegerAttr:
    return self.attributes["scatter_dim"]

def reduce_scatter(result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, scatter_dim: _Union[int, _ods_ir.IntegerAttr], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, reduction: _Optional[_Union[_Any, _ods_ir.Attribute]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult:
  return ReduceScatterOp(result=result, grid=grid, input=input, scatter_dim=scatter_dim, grid_axes=grid_axes, reduction=reduction, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class ScatterOp(_ods_ir.OpView):
  r"""
  For each device group defined by `grid_axes`, the input tensor on the `root`
  device is split along axis `scatter_dim` and distributed across the group.
  The content of the input on all other (non-root) devices is ignored.
  The `root` device is defined by its in-group multi-index.
  
  Example:
  ```
  shard.grid @grid0(shape = 2x2)
  %1 = shard.scatter %0 on @grid0 grid_axes = [0]
    scatter_dim = 0
    root = [1]
    : (tensor<2x2xi8>) -> tensor<1x2xi8>
  ```
  
  Input:
  ```
                            device
                            (0, 1)
                               ↓
                   +-------+-------+  | scatter tensor
  device (0, 0) -> |  *  * |  *  * |  | axis 0
                   |  *  * |  *  * |  ↓
                   +-------+-------+
  device (1, 0) -> |  1  2 |  5  6 |
                   |  3  4 |  7  8 |
                   +-------+-------+
                              ↑
                            device
                            (1, 1)
  ```
  
  Result:
  ```
                            device
                            (0, 1)
                               ↓
                   +-------+-------+
  device (0, 0) -> |  1  2 |  5  6 |
                   +-------+-------+ 
  device (1, 0) -> |  3  4 |  7  8 |
                   +-------+-------+
                              ↑
                            device
                            (1, 1)
  ```
  """

  OPERATION_NAME = "shard.scatter"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, scatter_dim: _Union[int, _ods_ir.IntegerAttr], root: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], root_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(input)
    operands.extend(_get_op_results_or_values(root_dynamic))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    if grid_axes is not None: attributes["grid_axes"] = (grid_axes if (
        isinstance(grid_axes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
          _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(grid_axes, context=_ods_context))
    attributes["scatter_dim"] = (scatter_dim if (
    isinstance(scatter_dim, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('IndexAttr')) else
      _ods_ir.AttrBuilder.get('IndexAttr')(scatter_dim, context=_ods_context))
    attributes["root"] = (root if (
    isinstance(root, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('DenseI64ArrayAttr')) else
      _ods_ir.AttrBuilder.get('DenseI64ArrayAttr')(root, context=_ods_context))
    results = []
    results.append(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operation.operands[0]

  @builtins.property
  def root_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operation.operands) - 2 + 1
    return self.operation.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["grid_axes"]

  @grid_axes.setter
  def grid_axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid_axes"] = value

  @builtins.property
  def scatter_dim(self) -> _ods_ir.IntegerAttr:
    return self.operation.attributes["scatter_dim"]

  @scatter_dim.setter
  def scatter_dim(self, value: _ods_ir.IntegerAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["scatter_dim"] = value

  @builtins.property
  def root(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.operation.attributes["root"]

  @root.setter
  def root(self, value: _ods_ir.DenseI64ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["root"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResult[_ods_ir.RankedTensorType]:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(ScatterOp)
class ScatterOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.scatter"

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operands[0]

  @builtins.property
  def root_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operands) - 2 + 1
    return self.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["grid_axes"]

  @builtins.property
  def scatter_dim(self) -> _ods_ir.IntegerAttr:
    return self.attributes["scatter_dim"]

  @builtins.property
  def root(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.attributes["root"]

def scatter(result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, scatter_dim: _Union[int, _ods_ir.IntegerAttr], root: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], root_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult[_ods_ir.RankedTensorType]:
  return ScatterOp(result=result, grid=grid, input=input, scatter_dim=scatter_dim, root=root, root_dynamic=root_dynamic, grid_axes=grid_axes, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class SendOp(_ods_ir.OpView):
  r"""
  Send input tensor to device `destination`, which is defined by its in-group
  multi-index. The groups are defined by `grid_axes`.
  """

  OPERATION_NAME = "shard.send"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, destination: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], destination_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(input)
    operands.extend(_get_op_results_or_values(destination_dynamic))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    if grid_axes is not None: attributes["grid_axes"] = (grid_axes if (
        isinstance(grid_axes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
          _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(grid_axes, context=_ods_context))
    attributes["destination"] = (destination if (
    isinstance(destination, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('DenseI64ArrayAttr')) else
      _ods_ir.AttrBuilder.get('DenseI64ArrayAttr')(destination, context=_ods_context))
    results = []
    results.append(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operation.operands[0]

  @builtins.property
  def destination_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operation.operands) - 2 + 1
    return self.operation.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["grid_axes"]

  @grid_axes.setter
  def grid_axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid_axes"] = value

  @builtins.property
  def destination(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.operation.attributes["destination"]

  @destination.setter
  def destination(self, value: _ods_ir.DenseI64ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["destination"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResult[_ods_ir.RankedTensorType]:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(SendOp)
class SendOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.send"

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operands[0]

  @builtins.property
  def destination_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operands) - 2 + 1
    return self.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["grid_axes"]

  @builtins.property
  def destination(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.attributes["destination"]

def send(result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, destination: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], destination_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult[_ods_ir.RankedTensorType]:
  return SendOp(result=result, grid=grid, input=input, destination=destination, destination_dynamic=destination_dynamic, grid_axes=grid_axes, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class ShardOp(_ods_ir.OpView):
  r"""
    The shard.shard operation is designed to specify and guide the sharding
    behavior of a tensor value across a grid topology. This operation has two
    operands and two optional attributes:
  
    1. `input`: This operand represents the tensor value that needs to be
    annotated for sharding.
  
    2. `sharding`: This attribute is type of `ShardingType`, which is the core data
    structure to represent distribution of a tensor on a shard. it is typically defined
    by an `shard.sharding` operation.
  
    3. `annotate_for_users`: A unit attribute addressing the scenario when a
    tensor's sharding annotation differs based on its context of use (either as
    a result or an operand). If specified, the sharding pertains to specific
    users of the tensor value, indicating how it should be considered when used
    as an operand in subsequent operations. If not, the sharding applies to the
    operation that defines the tensor value.
  
    Example:
    ```
  func.func @only_result_annotated(%arg0 : tensor<4x8xf32>) -> () {
      %sharding = shard.sharding @grid0 split_axes = [[0]] : !shard.sharding
      %0 = shard.shard %arg0 to %sharding : tensor<4x8xf32>
      ...
    }
  
    func.func @only_operand_annotated(%arg0 : tensor<4x8xf32>) -> () {
      %sharding = shard.sharding @grid0 split_axes = [[0]] : !shard.sharding
      %0 = shard.shard %arg0 to %sharding annotate_for_users : tensor<4x8xf32>
      ...
    }
    
    func.func @two_operands_annotated(%arg0 : tensor<4x8xf32>, %arg1 : tensor<16x8xf32>) -> () {
      %sharding = shard.sharding @grid0 split_axes = [[0]] : !shard.sharding
      %0 = shard.shard %arg0 to %sharding annotate_for_users : tensor<4x8xf32>
      %1 = shard.shard %arg1 to %sharding annotate_for_users : tensor<16x8xf32>
      ...
    }
  
    // The first shard.shard op applies to %arg0, the second shard.shard op
    // applies for the operand of op0, the third shard.shard op applies for the
    // operand of op2
    func.func @both_result_and_multi_operands_annotated(
        %arg0 : tensor<4x8xf32>) -> () {
      %sharding = shard.sharding @grid0 split_axes = [[0]] : !shard.sharding
      %0 = shard.shard %arg0 to %sharding : tensor<4x8xf32>
      %sharding1 = shard.sharding @grid0 split_axes = [[1]] : !shard.sharding
      %1 = shard.shard %0 to %sharding1 annotate_for_users : tensor<4x8xf32>
      %sharding2 = shard.sharding @grid0 split_axes = [[2]] : !shard.sharding
      %2 = shard.shard %0 to %sharding2 annotate_for_users : tensor<4x8xf32>
      "op0"(%1) : ...
      "op1"(%2) : ...
      ...
    }
    ```
  
    The following usages are undefined:
    ```
    func.func @annotate_on_same_result_with_different_sharding(
        %arg0 : tensor<4x8xf32>) -> () {
      %sharding1 = shard.sharding @grid0 split_axes = [[0]] : !shard.sharding
      %sharding2 = shard.sharding @grid0 split_axes = [[1]] : !shard.sharding
      %0 = shard.shard %arg0 to $sharding1 : tensor<4x8xf32>
      %1 = shard.shard %0 to sharding2 : tensor<4x8xf32>
      ...
    }
  
    func.func @annotate_on_same_result_same_value_with_different_sharding(
        %arg0 : tensor<4x8xf32>) -> () {
      %sharding1 = shard.sharding @grid0 split_axes = [[0]] : !shard.sharding
      %sharding2 = shard.sharding @grid0 split_axes = [[1]] : !shard.sharding
      %0 = shard.shard %arg0 to %sharding1 : tensor<4x8xf32>
      %1 = shard.shard %arg0 to %sharding2 : tensor<4x8xf32>
      ...
    }
  
    func.func @annotate_on_same_operand_with_different_sharding(
        %arg0 : tensor<4x8xf32>) -> () {
      %sharding1 = shard.sharding @grid0 split_axes = [[0]] : !shard.sharding
      %sharding2 = shard.sharding @grid0 split_axes = [[1]] : !shard.sharding
      %0 = shard.shard %arg0 to %sharding1 annotate_for_users : tensor<4x8xf32>
      %1 = shard.shard %0 to %sharding2 annotate_for_users : tensor<4x8xf32>
      ...
    }
  
    func.func @result_annotated_after_operand(
        %arg0 : tensor<4x8xf32>) -> () {
      %sharding1 = shard.sharding @grid0 split_axes = [[0]] : !shard.sharding
      %sharding2 = shard.sharding @grid0 split_axes = [[1]] : !shard.sharding
      %0 = shard.shard %arg0 to %sharding1 annotate_for_users : tensor<4x8xf32>
      %1 = shard.shard %0 to %sharding2 : tensor<4x8xf32>
      ...
    }
    ```
  """

  OPERATION_NAME = "shard.shard"

  _ODS_REGIONS = (0, True)

  def __init__(self, src: _ods_ir.Value[_ods_ir.RankedTensorType], sharding: _ods_ir.Value, *, annotate_for_users: _Optional[bool] = None, results: _Optional[_Sequence[_ods_ir.Type]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(src)
    operands.append(sharding)
    _ods_context = _ods_get_default_loc_context(loc)
    if bool(annotate_for_users): attributes["annotate_for_users"] = _ods_ir.UnitAttr.get(
      _ods_get_default_loc_context(loc))
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def src(self) -> _ods_ir.Value[_ods_ir.RankedTensorType]:
    return self.operation.operands[0]

  @builtins.property
  def sharding(self) -> _ods_ir.Value:
    return self.operation.operands[1]

  @builtins.property
  def annotate_for_users(self) -> bool:
    return "annotate_for_users" in self.operation.attributes

  @annotate_for_users.setter
  def annotate_for_users(self, value):
    if bool(value):
      self.operation.attributes["annotate_for_users"] = _ods_ir.UnitAttr.get()
    elif "annotate_for_users" in self.operation.attributes:
      del self.operation.attributes["annotate_for_users"]

  @annotate_for_users.deleter
  def annotate_for_users(self):
    del self.operation.attributes["annotate_for_users"]

  @builtins.property
  def result(self) -> _ods_ir.OpResult[_ods_ir.RankedTensorType]:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(ShardOp)
class ShardOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.shard"

  @builtins.property
  def src(self) -> _ods_ir.Value[_ods_ir.RankedTensorType]:
    return self.operands[0]

  @builtins.property
  def sharding(self) -> _ods_ir.Value:
    return self.operands[1]

  @builtins.property
  def annotate_for_users(self) -> bool:
    return "annotate_for_users" in self.attributes

def shard(src: _ods_ir.Value[_ods_ir.RankedTensorType], sharding: _ods_ir.Value, *, annotate_for_users: _Optional[bool] = None, results: _Optional[_Sequence[_ods_ir.Type]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult[_ods_ir.RankedTensorType]:
  return ShardOp(src=src, sharding=sharding, annotate_for_users=annotate_for_users, results=results, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class ShardShapeOp(_ods_ir.OpView):
  r"""
  The device/process id is a multi-index of the device/process in the shard.
  This operation might be used during partition when the shard shape depends
  on (non-constant) values used in `shard.sharding`.
  """

  OPERATION_NAME = "shard.shard_shape"

  _ODS_OPERAND_SEGMENTS = [-1,1,-1,]

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _Sequence[_ods_ir.Type], dims: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], dims_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], sharding: _ods_ir.Value, device: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], device_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], *, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(_get_op_results_or_values(dims_dynamic))
    operands.append(sharding)
    operands.append(_get_op_results_or_values(device_dynamic))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["dims"] = (dims if (
    isinstance(dims, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('DenseI64ArrayAttr')) else
      _ods_ir.AttrBuilder.get('DenseI64ArrayAttr')(dims, context=_ods_context))
    attributes["device"] = (device if (
    isinstance(device, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('DenseI64ArrayAttr')) else
      _ods_ir.AttrBuilder.get('DenseI64ArrayAttr')(device, context=_ods_context))
    results = []
    results.extend(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def dims_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    operand_range = _ods_segmented_accessor(
         self.operation.operands,
         self.operation.attributes["operandSegmentSizes"], 0)
    return operand_range

  @builtins.property
  def sharding(self) -> _ods_ir.Value:
    operand_range = _ods_segmented_accessor(
         self.operation.operands,
         self.operation.attributes["operandSegmentSizes"], 1)
    return operand_range[0]

  @builtins.property
  def device_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    operand_range = _ods_segmented_accessor(
         self.operation.operands,
         self.operation.attributes["operandSegmentSizes"], 2)
    return operand_range

  @builtins.property
  def dims(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.operation.attributes["dims"]

  @dims.setter
  def dims(self, value: _ods_ir.DenseI64ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["dims"] = value

  @builtins.property
  def device(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.operation.attributes["device"]

  @device.setter
  def device(self, value: _ods_ir.DenseI64ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["device"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResultList[_ods_ir.IndexType]:
    _ods_variadic_group_length = len(self.operation.results) - 1 + 1
    return self.operation.results[0:0 + _ods_variadic_group_length]

@_ods_cext.register_op_adaptor(ShardShapeOp)
class ShardShapeOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.shard_shape"

  @builtins.property
  def dims_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    operand_range = _ods_segmented_accessor(
         self.operands,
         self.attributes["operandSegmentSizes"], 0)
    return operand_range

  @builtins.property
  def sharding(self) -> _ods_ir.Value:
    operand_range = _ods_segmented_accessor(
         self.operands,
         self.attributes["operandSegmentSizes"], 1)
    return operand_range[0]

  @builtins.property
  def device_dynamic(self) -> _ods_ir.OpOperandList[_ods_ir.IndexType]:
    operand_range = _ods_segmented_accessor(
         self.operands,
         self.attributes["operandSegmentSizes"], 2)
    return operand_range

  @builtins.property
  def dims(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.attributes["dims"]

  @builtins.property
  def device(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.attributes["device"]

def shard_shape(result: _Sequence[_ods_ir.Type], dims: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], dims_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], sharding: _ods_ir.Value, device: _Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr], device_dynamic: _Sequence[_ods_ir.Value[_ods_ir.IndexType]], *, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _Union[_ods_ir.OpResult, _ods_ir.OpResultList, ShardShapeOp]:
  op = ShardShapeOp(result=result, dims=dims, dims_dynamic=dims_dynamic, sharding=sharding, device=device, device_dynamic=device_dynamic, loc=loc, ip=ip); results = op.results
  return results if len(results) > 1 else (results[0] if len(results) == 1 else op)

@_ods_cext.register_operation(_Dialect)
class ShardingOp(_ods_ir.OpView):
  r"""
  The Sharding specifies how a tensor is sharded and distributed across the
  process shard. It is typically used in a `shard.shard` operation.
  The operation has the following attributes and operands:
  
  1. `grid`: this attribute is a FlatSymbolRefAttr that refers to the device
  grid where the distributed tensor is placed. The symbol must resolve to a
  `shard.grid` operation.
  
  2. `split_axes`: is an array composed of int64_t sub-arrays. The outer array's
  maximum size is the `rank` of the related tensor. For the i-th sub-array, if
  its value is [x, y], it indicates that the tensor's i-th dimension is splitted
  along the x and y axes of the device grid.
  
  3. [Optional] Sizes of halos to be added for each sharded tensor dimension.
  `halo_sizes` is provided as a flattened 1d array of i64s, 2 values for each
  sharded dimension. `halo_sizes = [1, 2]` means that the first sharded dimension
  gets an additional halo of size 1 at the start of the first dimension and a halo
  size is 2 at its end. `halo_sizes = [1, 2, 2, 3]` defines halos for the first 2
  sharded dimensions e.g. the first sharded dimension gets `[1,2]` halos and the
  seconds gets `[2,3]` halos. `?` indicates dynamic halo sizes.
  
  4. [Optional] Offsets for each shard and sharded tensor dimension.
  `sharded_dims_offsets` is provided as a flattened 1d array of i64s. For each
  sharded tensor dimension the offsets (starting index) of all shards in that
  dimension and an additional value for the end of the last shard are provided.
  For a 1d sharding this means that position `i` has the exclusive prefix sum for
  shard `i`, and since only contiguous sharding is supported, its inclusive prefix
  sum is at position 'i+1'.
  
  Assuming a 3d-tensor of shape 32x32x32 with the first 2 dimensions being sharded,
  `sharded_dims_offsets` = [0, 24, 32, 0, 20, 32] means that the first device of
  the device-grid will get a shard of shape 24x20x32 and the second device will get
  a shard of shape 8x12x32. `?` indicates dynamic shard dimensions.
  
  `halo_sizes` and `sharded_dims_offsets` are mutually exclusive.
  
  Examples:
  
  ```
  shard.grid @grid0(shape = 2x2x4)
  shard.grid @grid1d_4(shape = 4)
  
  // The tensor is fully replicated on @grid0.
  // Currently, there must be at least one sub-array present in axes, even
  // if it's empty. Otherwise, a parsing error will occur.
  %sharding0 = shard.sharding @grid0 split_axes = [[]]
  
  // The tensor is sharded on the first dimension along axis 0 of @grid0
  %sharding1 = shard.sharding @grid0 split_axes = [[0]]
  
  // Could be used for a shard.shard op
  %sharded0 = shard.shard %arg0 to %sharding3 : tensor<4x8xf32>
  
  // The tensor is sharded on its first dimension along axis 0 of @grid0 and
  // and it has halo-sizes of 1 and 2 on the sharded dim.
  %halo_sharding = shard.sharding @grid0 split_axes = [[0]] halo_sizes = [1, 2]
  %sharded1 = shard.shard %arg0 to %halo_sharding : tensor<4x8xf32>
  
  // The tensor is sharded on its second dimension along axis 0 of @grid1d_4
  // and it has pre-defined shard sizes. The shards of the devices will have
  // the following shapes: [4x2, 4x3, 4x4, 4x5]
  %sharding4 = shard.sharding @grid1d_4 split_axes = [[], [0]] sharded_dims_offsets = [0, 2, 5, 9, 14]
  %sharded2 = shard.shard %arg0 to %sharding4 : tensor<4x14xf32>
  ```
  """

  OPERATION_NAME = "shard.sharding"

  _ODS_OPERAND_SEGMENTS = [-1,-1,]

  _ODS_REGIONS = (0, True)

  def __init__(self, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], split_axes: _Union[_Any, _ods_ir.Attribute], dynamic_sharded_dims_offsets: _Sequence[_ods_ir.Value[_ods_ir.IntegerType]], dynamic_halo_sizes: _Sequence[_ods_ir.Value[_ods_ir.IntegerType]], *, static_sharded_dims_offsets: _Optional[_Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr]] = None, static_halo_sizes: _Optional[_Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr]] = None, results: _Optional[_Sequence[_ods_ir.Type]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(_get_op_results_or_values(dynamic_sharded_dims_offsets))
    operands.append(_get_op_results_or_values(dynamic_halo_sizes))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    attributes["split_axes"] = (split_axes if (
    isinstance(split_axes, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('Shard_GridAxesArrayAttr')) else
      _ods_ir.AttrBuilder.get('Shard_GridAxesArrayAttr')(split_axes, context=_ods_context))
    if static_sharded_dims_offsets is not None: attributes["static_sharded_dims_offsets"] = (static_sharded_dims_offsets if (
        isinstance(static_sharded_dims_offsets, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('DenseI64ArrayAttr')) else
          _ods_ir.AttrBuilder.get('DenseI64ArrayAttr')(static_sharded_dims_offsets, context=_ods_context))
    if static_halo_sizes is not None: attributes["static_halo_sizes"] = (static_halo_sizes if (
        isinstance(static_halo_sizes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('DenseI64ArrayAttr')) else
          _ods_ir.AttrBuilder.get('DenseI64ArrayAttr')(static_halo_sizes, context=_ods_context))
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def dynamic_sharded_dims_offsets(self) -> _ods_ir.OpOperandList[_ods_ir.IntegerType]:
    operand_range = _ods_segmented_accessor(
         self.operation.operands,
         self.operation.attributes["operandSegmentSizes"], 0)
    return operand_range

  @builtins.property
  def dynamic_halo_sizes(self) -> _ods_ir.OpOperandList[_ods_ir.IntegerType]:
    operand_range = _ods_segmented_accessor(
         self.operation.operands,
         self.operation.attributes["operandSegmentSizes"], 1)
    return operand_range

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def split_axes(self) -> _ods_ir.Attribute:
    return self.operation.attributes["split_axes"]

  @split_axes.setter
  def split_axes(self, value: _ods_ir.Attribute):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["split_axes"] = value

  @builtins.property
  def static_sharded_dims_offsets(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.operation.attributes["static_sharded_dims_offsets"]

  @static_sharded_dims_offsets.setter
  def static_sharded_dims_offsets(self, value: _ods_ir.DenseI64ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["static_sharded_dims_offsets"] = value

  @builtins.property
  def static_halo_sizes(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.operation.attributes["static_halo_sizes"]

  @static_halo_sizes.setter
  def static_halo_sizes(self, value: _ods_ir.DenseI64ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["static_halo_sizes"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResult:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(ShardingOp)
class ShardingOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.sharding"

  @builtins.property
  def dynamic_sharded_dims_offsets(self) -> _ods_ir.OpOperandList[_ods_ir.IntegerType]:
    operand_range = _ods_segmented_accessor(
         self.operands,
         self.attributes["operandSegmentSizes"], 0)
    return operand_range

  @builtins.property
  def dynamic_halo_sizes(self) -> _ods_ir.OpOperandList[_ods_ir.IntegerType]:
    operand_range = _ods_segmented_accessor(
         self.operands,
         self.attributes["operandSegmentSizes"], 1)
    return operand_range

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def split_axes(self) -> _ods_ir.Attribute:
    return self.attributes["split_axes"]

  @builtins.property
  def static_sharded_dims_offsets(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.attributes["static_sharded_dims_offsets"]

  @builtins.property
  def static_halo_sizes(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.attributes["static_halo_sizes"]

def sharding(grid: _Union[str, _ods_ir.FlatSymbolRefAttr], split_axes: _Union[_Any, _ods_ir.Attribute], dynamic_sharded_dims_offsets: _Sequence[_ods_ir.Value[_ods_ir.IntegerType]], dynamic_halo_sizes: _Sequence[_ods_ir.Value[_ods_ir.IntegerType]], *, static_sharded_dims_offsets: _Optional[_Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr]] = None, static_halo_sizes: _Optional[_Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr]] = None, results: _Optional[_Sequence[_ods_ir.Type]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult:
  return ShardingOp(grid=grid, split_axes=split_axes, dynamic_sharded_dims_offsets=dynamic_sharded_dims_offsets, dynamic_halo_sizes=dynamic_halo_sizes, static_sharded_dims_offsets=static_sharded_dims_offsets, static_halo_sizes=static_halo_sizes, results=results, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class ShiftOp(_ods_ir.OpView):
  r"""
  Within each device group defined by `grid_axes`, shifts input tensors along the
  device grid's axis `shift_axis` by the specified offset. The `shift_axis` must
  be one of the `grid_axes`. If the `rotate` attribute is set, the shift is circular.
  That is, the offset wraps around according to the group size along `shift_axis`.
  Otherwise, the results on devices without a corresponding source are undefined.
  
  Example:
  ```
  shard.grid @grid0(shape = 2x4)
  %1 = shard.shift on @grid0 grid_axes = [1]
    shift_axis = 1 offset = 2 rotate
    : tensor<2xi8> -> tensor<2xi8>
  ```
  
  Input:
  ```
  grid axis 1
  ----------->
  
  +----+----+----+----+
  |  1 |  2 |  3 |  4 |
  +----+----+----+----+
  |  5 |  6 |  7 |  8 |
  +----+----+----+----+
  ```
  
  Result:
  ```
  +----+----+----+----+
  |  3 |  4 |  1 |  2 |
  +----+----+----+----+
  |  7 |  8 |  5 |  6 |
  +----+----+----+----+
  ```
  """

  OPERATION_NAME = "shard.shift"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, shift_axis: _Union[int, _ods_ir.IntegerAttr], offset: _Union[int, _ods_ir.IntegerAttr], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, rotate: _Optional[bool] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(input)
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    if grid_axes is not None: attributes["grid_axes"] = (grid_axes if (
        isinstance(grid_axes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('Shard_GridAxesAttr')) else
          _ods_ir.AttrBuilder.get('Shard_GridAxesAttr')(grid_axes, context=_ods_context))
    attributes["shift_axis"] = (shift_axis if (
    isinstance(shift_axis, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('IndexAttr')) else
      _ods_ir.AttrBuilder.get('IndexAttr')(shift_axis, context=_ods_context))
    attributes["offset"] = (offset if (
    isinstance(offset, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I64Attr')) else
      _ods_ir.AttrBuilder.get('I64Attr')(offset, context=_ods_context))
    if bool(rotate): attributes["rotate"] = _ods_ir.UnitAttr.get(
      _ods_get_default_loc_context(loc))
    results = []
    results.append(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operation.operands[0]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.operation.attributes["grid_axes"]

  @grid_axes.setter
  def grid_axes(self, value: _ods_ir.DenseI16ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid_axes"] = value

  @builtins.property
  def shift_axis(self) -> _ods_ir.IntegerAttr:
    return self.operation.attributes["shift_axis"]

  @shift_axis.setter
  def shift_axis(self, value: _ods_ir.IntegerAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["shift_axis"] = value

  @builtins.property
  def offset(self) -> _ods_ir.IntegerAttr:
    return self.operation.attributes["offset"]

  @offset.setter
  def offset(self, value: _ods_ir.IntegerAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["offset"] = value

  @builtins.property
  def rotate(self) -> bool:
    return "rotate" in self.operation.attributes

  @rotate.setter
  def rotate(self, value):
    if bool(value):
      self.operation.attributes["rotate"] = _ods_ir.UnitAttr.get()
    elif "rotate" in self.operation.attributes:
      del self.operation.attributes["rotate"]

  @rotate.deleter
  def rotate(self):
    del self.operation.attributes["rotate"]

  @builtins.property
  def result(self) -> _ods_ir.OpResult[_ods_ir.RankedTensorType]:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(ShiftOp)
class ShiftOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.shift"

  @builtins.property
  def input(self) -> _ods_ir.Value:
    return self.operands[0]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def grid_axes(self) -> _ods_ir.DenseI16ArrayAttr:
    return self.attributes["grid_axes"]

  @builtins.property
  def shift_axis(self) -> _ods_ir.IntegerAttr:
    return self.attributes["shift_axis"]

  @builtins.property
  def offset(self) -> _ods_ir.IntegerAttr:
    return self.attributes["offset"]

  @builtins.property
  def rotate(self) -> bool:
    return "rotate" in self.attributes

def shift(result: _ods_ir.Type, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], input: _ods_ir.Value, shift_axis: _Union[int, _ods_ir.IntegerAttr], offset: _Union[int, _ods_ir.IntegerAttr], *, grid_axes: _Optional[_Union[_Any, _ods_ir.DenseI16ArrayAttr]] = None, rotate: _Optional[bool] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult[_ods_ir.RankedTensorType]:
  return ShiftOp(result=result, grid=grid, input=input, shift_axis=shift_axis, offset=offset, grid_axes=grid_axes, rotate=rotate, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class UpdateHaloOp(_ods_ir.OpView):
  r"""
  This operation updates halo regions of shards, e.g. if their sharding
  specified halos and the actual tensor/memref data might have changed
  on the remote devices. Changes might be caused by mutating operations
  and/or if the new halo regions are larger than the existing ones.
  
  Destination is supposed to be initialized with the local data (not halos).
  
  Assumes all devices hold tensors with same-sized halo data as specified
  by `source_halo_sizes/static_source_halo_sizes` and
  `destination_halo_sizes/static_destination_halo_sizes` in source shard
  and destination/result shard.
  
  `split_axes` specifies for each tensor axis along which grid axes its halo
  data is updated.
  
  """

  OPERATION_NAME = "shard.update_halo"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, destination: _ods_ir.Value, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], split_axes: _Union[_Any, _ods_ir.Attribute], halo_sizes: _Sequence[_ods_ir.Value[_ods_ir.IntegerType]], *, static_halo_sizes: _Optional[_Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(destination)
    operands.extend(_get_op_results_or_values(halo_sizes))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["grid"] = (grid if (
    isinstance(grid, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
      _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(grid, context=_ods_context))
    attributes["split_axes"] = (split_axes if (
    isinstance(split_axes, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('Shard_GridAxesArrayAttr')) else
      _ods_ir.AttrBuilder.get('Shard_GridAxesArrayAttr')(split_axes, context=_ods_context))
    if static_halo_sizes is not None: attributes["static_halo_sizes"] = (static_halo_sizes if (
        isinstance(static_halo_sizes, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('DenseI64ArrayAttr')) else
          _ods_ir.AttrBuilder.get('DenseI64ArrayAttr')(static_halo_sizes, context=_ods_context))
    results = []
    results.append(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def destination(self) -> _ods_ir.Value:
    return self.operation.operands[0]

  @builtins.property
  def halo_sizes(self) -> _ods_ir.OpOperandList[_ods_ir.IntegerType]:
    _ods_variadic_group_length = len(self.operation.operands) - 2 + 1
    return self.operation.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.operation.attributes["grid"]

  @grid.setter
  def grid(self, value: _ods_ir.FlatSymbolRefAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["grid"] = value

  @builtins.property
  def split_axes(self) -> _ods_ir.Attribute:
    return self.operation.attributes["split_axes"]

  @split_axes.setter
  def split_axes(self, value: _ods_ir.Attribute):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["split_axes"] = value

  @builtins.property
  def static_halo_sizes(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.operation.attributes["static_halo_sizes"]

  @static_halo_sizes.setter
  def static_halo_sizes(self, value: _ods_ir.DenseI64ArrayAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["static_halo_sizes"] = value

  @builtins.property
  def result(self) -> _ods_ir.OpResult:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(UpdateHaloOp)
class UpdateHaloOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "shard.update_halo"

  @builtins.property
  def destination(self) -> _ods_ir.Value:
    return self.operands[0]

  @builtins.property
  def halo_sizes(self) -> _ods_ir.OpOperandList[_ods_ir.IntegerType]:
    _ods_variadic_group_length = len(self.operands) - 2 + 1
    return self.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def grid(self) -> _ods_ir.FlatSymbolRefAttr:
    return self.attributes["grid"]

  @builtins.property
  def split_axes(self) -> _ods_ir.Attribute:
    return self.attributes["split_axes"]

  @builtins.property
  def static_halo_sizes(self) -> _ods_ir.DenseI64ArrayAttr:
    return self.attributes["static_halo_sizes"]

def update_halo(result: _ods_ir.Type, destination: _ods_ir.Value, grid: _Union[str, _ods_ir.FlatSymbolRefAttr], split_axes: _Union[_Any, _ods_ir.Attribute], halo_sizes: _Sequence[_ods_ir.Value[_ods_ir.IntegerType]], *, static_halo_sizes: _Optional[_Union[_Sequence[int], _ods_ir.DenseI64ArrayAttr]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult:
  return UpdateHaloOp(result=result, destination=destination, grid=grid, split_axes=split_axes, halo_sizes=halo_sizes, static_halo_sizes=static_halo_sizes, loc=loc, ip=ip).result
