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

from jaxlib.mlir.dialects._ods_common import _cext as _ods_cext
from jaxlib.mlir.dialects._ods_common import extend_opview_class as _ods_extend_opview_class, segmented_accessor as _ods_segmented_accessor, equally_sized_accessor as _ods_equally_sized_accessor, get_default_loc_context as _ods_get_default_loc_context, get_op_result_or_value as _get_op_result_or_value, get_op_results_or_values as _get_op_results_or_values
_ods_ir = _ods_cext.ir

try:
  from . import _tpu_ops_ext as _ods_ext_module
except ImportError:
  _ods_ext_module = None

import builtins


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


@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class AllReduceOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.all_reduce"

  _ODS_REGIONS = (0, True)

  def __init__(self, input, dim, kind, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(input))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["dim"] = (dim if (
    issubclass(type(dim), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I64Attr')) else
      _ods_ir.AttrBuilder.get('I64Attr')(dim, context=_ods_context))
    attributes["kind"] = (kind if (
    issubclass(type(kind), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('TPU_ReductionKindAttr')) else
      _ods_ir.AttrBuilder.get('TPU_ReductionKindAttr')(kind, context=_ods_context))
    results.extend([operands[0].type] * 1)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

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

  @builtins.property
  def dim(self):
    return self.operation.attributes["dim"]

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

  @builtins.property
  def kind(self):
    return self.operation.attributes["kind"]

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

  @builtins.property
  def output(self):
    return self.operation.results[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class AllocaSemaphoreOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.sem_alloc"

  _ODS_REGIONS = (0, True)

  def __init__(self, result, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    results.append(result)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

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

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class BitcastOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.bitcast"

  _ODS_REGIONS = (0, True)

  def __init__(self, output, input, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(input))
    _ods_context = _ods_get_default_loc_context(loc)
    results.append(output)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

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

  @builtins.property
  def output(self):
    return self.operation.results[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class BroadcastInSublanesOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.broadcast_in_sublanes"

  _ODS_REGIONS = (0, True)

  def __init__(self, output, source, lane, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(source))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["lane"] = (lane if (
    issubclass(type(lane), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I32Attr')) else
      _ods_ir.AttrBuilder.get('I32Attr')(lane, context=_ods_context))
    results.append(output)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def source(self):
    return self.operation.operands[0]

  @builtins.property
  def lane(self):
    return self.operation.attributes["lane"]

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

  @builtins.property
  def output(self):
    return self.operation.results[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class ConcatenateOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.concatenate"

  _ODS_REGIONS = (0, True)

  def __init__(self, output, sources, dimension, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.extend(_get_op_results_or_values(sources))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["dimension"] = (dimension if (
    issubclass(type(dimension), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I32Attr')) else
      _ods_ir.AttrBuilder.get('I32Attr')(dimension, context=_ods_context))
    results.append(output)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def sources(self):
    _ods_variadic_group_length = len(self.operation.operands) - 1 + 1
    return self.operation.operands[0:0 + _ods_variadic_group_length]

  @builtins.property
  def dimension(self):
    return self.operation.attributes["dimension"]

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

  @builtins.property
  def output(self):
    return self.operation.results[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class CreateMaskOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.create_mask"

  _ODS_REGIONS = (0, True)

  def __init__(self, output, low, high, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.extend(_get_op_results_or_values(low))
    operands.extend(_get_op_results_or_values(high))
    _ods_context = _ods_get_default_loc_context(loc)
    results.append(output)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def low(self):
    start, pg = _ods_equally_sized_accessor(operation.operands, 2, 0, 0)
    return self.operation.operands[start:start + pg]

  @builtins.property
  def high(self):
    start, pg = _ods_equally_sized_accessor(operation.operands, 2, 0, 1)
    return self.operation.operands[start:start + pg]

  @builtins.property
  def output(self):
    return self.operation.results[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class CreateSubelementMaskOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.create_subelement_mask"

  _ODS_REGIONS = (0, True)

  def __init__(self, output, from_, to, num_subelems, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["from"] = (from_ if (
    issubclass(type(from_), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I32Attr')) else
      _ods_ir.AttrBuilder.get('I32Attr')(from_, context=_ods_context))
    attributes["to"] = (to if (
    issubclass(type(to), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I32Attr')) else
      _ods_ir.AttrBuilder.get('I32Attr')(to, context=_ods_context))
    attributes["num_subelems"] = (num_subelems if (
    issubclass(type(num_subelems), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I32Attr')) else
      _ods_ir.AttrBuilder.get('I32Attr')(num_subelems, context=_ods_context))
    results.append(output)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def from_(self):
    return self.operation.attributes["from"]

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

  @builtins.property
  def to(self):
    return self.operation.attributes["to"]

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

  @builtins.property
  def num_subelems(self):
    return self.operation.attributes["num_subelems"]

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

  @builtins.property
  def output(self):
    return self.operation.results[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class DeviceIdOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.device_id"

  _ODS_REGIONS = (0, True)

  def __init__(self, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

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

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class DynamicGatherOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.dynamic_gather"

  _ODS_REGIONS = (0, True)

  def __init__(self, output, source, indices, dimension, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(source))
    operands.append(_get_op_result_or_value(indices))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["dimension"] = (dimension if (
    issubclass(type(dimension), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I32Attr')) else
      _ods_ir.AttrBuilder.get('I32Attr')(dimension, context=_ods_context))
    results.append(output)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def source(self):
    return self.operation.operands[0]

  @builtins.property
  def indices(self):
    return self.operation.operands[1]

  @builtins.property
  def dimension(self):
    return self.operation.attributes["dimension"]

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

  @builtins.property
  def output(self):
    return self.operation.results[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class EnqueueDMAOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.enqueue_dma"

  _ODS_REGIONS = (0, True)

  def __init__(self, source, target, target_semaphore, *, source_semaphore=None, device_id=None, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(source))
    if source_semaphore is not None: operands.append(_get_op_result_or_value(source_semaphore))
    operands.append(_get_op_result_or_value(target))
    operands.append(_get_op_result_or_value(target_semaphore))
    if device_id is not None: operands.append(_get_op_result_or_value(device_id))
    _ods_context = _ods_get_default_loc_context(loc)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def source(self):
    start, pg = _ods_equally_sized_accessor(operation.operands, 2, 0, 0)
    return self.operation.operands[start]

  @builtins.property
  def source_semaphore(self):
    start, pg = _ods_equally_sized_accessor(operation.operands, 2, 1, 0)
    return self.operation.operands[start:start + pg]

  @builtins.property
  def target(self):
    start, pg = _ods_equally_sized_accessor(operation.operands, 2, 1, 1)
    return self.operation.operands[start]

  @builtins.property
  def target_semaphore(self):
    start, pg = _ods_equally_sized_accessor(operation.operands, 2, 2, 1)
    return self.operation.operands[start]

  @builtins.property
  def device_id(self):
    start, pg = _ods_equally_sized_accessor(operation.operands, 2, 3, 1)
    return self.operation.operands[start:start + pg]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class EraseLayoutOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.erase_memref_layout"

  _ODS_REGIONS = (0, True)

  def __init__(self, result, operand, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(operand))
    _ods_context = _ods_get_default_loc_context(loc)
    results.append(result)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def operand(self):
    return self.operation.operands[0]

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

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class GatherOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.gather"

  _ODS_REGIONS = (0, True)

  def __init__(self, output, source, indices, dimension, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(source))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["indices"] = (indices if (
    issubclass(type(indices), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('DenseI32ArrayAttr')) else
      _ods_ir.AttrBuilder.get('DenseI32ArrayAttr')(indices, context=_ods_context))
    attributes["dimension"] = (dimension if (
    issubclass(type(dimension), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I32Attr')) else
      _ods_ir.AttrBuilder.get('I32Attr')(dimension, context=_ods_context))
    results.append(output)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def source(self):
    return self.operation.operands[0]

  @builtins.property
  def indices(self):
    return self.operation.attributes["indices"]

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

  @builtins.property
  def dimension(self):
    return self.operation.attributes["dimension"]

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

  @builtins.property
  def output(self):
    return self.operation.results[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class GetBarrierSemaphoreOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.sem_barrier"

  _ODS_REGIONS = (0, True)

  def __init__(self, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def semaphore(self):
    return self.operation.results[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class IotaOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.iota"

  _ODS_REGIONS = (0, True)

  def __init__(self, output, *, dimension=None, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    if dimension is not None: attributes["dimension"] = (dimension if (
        issubclass(type(dimension), _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('I32Attr')) else
          _ods_ir.AttrBuilder.get('I32Attr')(dimension, context=_ods_context))
    results.append(output)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def dimension(self):
    if "dimension" not in self.operation.attributes:
      return None
    return self.operation.attributes["dimension"]

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

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

  @builtins.property
  def output(self):
    return self.operation.results[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class LoadOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.load"

  _ODS_REGIONS = (0, True)

  def __init__(self, result, base, indices, sublane_mask, *, sublane_stride=None, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(base))
    operands.extend(_get_op_results_or_values(indices))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["sublane_mask"] = (sublane_mask if (
    issubclass(type(sublane_mask), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('DenseBoolArrayAttr')) else
      _ods_ir.AttrBuilder.get('DenseBoolArrayAttr')(sublane_mask, context=_ods_context))
    if sublane_stride is not None: attributes["sublane_stride"] = (sublane_stride if (
        issubclass(type(sublane_stride), _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('I32Attr')) else
          _ods_ir.AttrBuilder.get('I32Attr')(sublane_stride, context=_ods_context))
    results.append(result)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def base(self):
    return self.operation.operands[0]

  @builtins.property
  def indices(self):
    _ods_variadic_group_length = len(self.operation.operands) - 2 + 1
    return self.operation.operands[1:1 + _ods_variadic_group_length]

  @builtins.property
  def sublane_mask(self):
    return self.operation.attributes["sublane_mask"]

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

  @builtins.property
  def sublane_stride(self):
    if "sublane_stride" not in self.operation.attributes:
      return None
    return self.operation.attributes["sublane_stride"]

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

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

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

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class MatmulOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.matmul"

  _ODS_REGIONS = (0, True)

  def __init__(self, result, lhs, rhs, acc, *, transpose_lhs=None, transpose_rhs=None, precision=None, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(lhs))
    operands.append(_get_op_result_or_value(rhs))
    operands.append(_get_op_result_or_value(acc))
    _ods_context = _ods_get_default_loc_context(loc)
    if transpose_lhs is not None: attributes["transpose_lhs"] = (transpose_lhs if (
        issubclass(type(transpose_lhs), _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('BoolAttr')) else
          _ods_ir.AttrBuilder.get('BoolAttr')(transpose_lhs, context=_ods_context))
    if transpose_rhs is not None: attributes["transpose_rhs"] = (transpose_rhs if (
        issubclass(type(transpose_rhs), _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('BoolAttr')) else
          _ods_ir.AttrBuilder.get('BoolAttr')(transpose_rhs, context=_ods_context))
    if precision is not None: attributes["precision"] = (precision if (
        issubclass(type(precision), _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('TPU_ContractPrecisionEnum')) else
          _ods_ir.AttrBuilder.get('TPU_ContractPrecisionEnum')(precision, context=_ods_context))
    results.append(result)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def lhs(self):
    return self.operation.operands[0]

  @builtins.property
  def rhs(self):
    return self.operation.operands[1]

  @builtins.property
  def acc(self):
    return self.operation.operands[2]

  @builtins.property
  def transpose_lhs(self):
    return self.operation.attributes["transpose_lhs"]

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

  @builtins.property
  def transpose_rhs(self):
    return self.operation.attributes["transpose_rhs"]

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

  @builtins.property
  def precision(self):
    if "precision" not in self.operation.attributes:
      return None
    return self.operation.attributes["precision"]

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

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

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

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class MemRefSliceOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.memref_slice"

  _ODS_REGIONS = (0, True)

  def __init__(self, result, mem_ref, base_idx, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(mem_ref))
    operands.extend(_get_op_results_or_values(base_idx))
    _ods_context = _ods_get_default_loc_context(loc)
    results.append(result)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def mem_ref(self):
    return self.operation.operands[0]

  @builtins.property
  def base_idx(self):
    _ods_variadic_group_length = len(self.operation.operands) - 2 + 1
    return self.operation.operands[1:1 + _ods_variadic_group_length]

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

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class PackSubelementsOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.pack_subelements"

  _ODS_REGIONS = (0, True)

  def __init__(self, output, sources, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.extend(_get_op_results_or_values(sources))
    _ods_context = _ods_get_default_loc_context(loc)
    results.append(output)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def sources(self):
    _ods_variadic_group_length = len(self.operation.operands) - 1 + 1
    return self.operation.operands[0:0 + _ods_variadic_group_length]

  @builtins.property
  def output(self):
    return self.operation.results[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class RegionOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.region"

  _ODS_REGIONS = (1, True)

  def __init__(self, results_, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    results.extend(results_)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def results_(self):
    _ods_variadic_group_length = len(self.operation.results) - 1 + 1
    return self.operation.results[0:0 + _ods_variadic_group_length]

  @builtins.property
  def region(self):
    return self.regions[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class ReinterpretCastOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.reinterpret_cast"

  _ODS_REGIONS = (0, True)

  def __init__(self, result, input, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(input))
    _ods_context = _ods_get_default_loc_context(loc)
    results.append(result)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

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

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

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class RepeatOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.repeat"

  _ODS_REGIONS = (0, True)

  def __init__(self, output, source, dimension, times, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(source))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["dimension"] = (dimension if (
    issubclass(type(dimension), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I32Attr')) else
      _ods_ir.AttrBuilder.get('I32Attr')(dimension, context=_ods_context))
    attributes["times"] = (times if (
    issubclass(type(times), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I32Attr')) else
      _ods_ir.AttrBuilder.get('I32Attr')(times, context=_ods_context))
    results.append(output)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def source(self):
    return self.operation.operands[0]

  @builtins.property
  def dimension(self):
    return self.operation.attributes["dimension"]

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

  @builtins.property
  def times(self):
    return self.operation.attributes["times"]

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

  @builtins.property
  def output(self):
    return self.operation.results[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class RollVectorsOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.roll_vectors"

  _ODS_REGIONS = (0, True)

  def __init__(self, output, input, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.extend(_get_op_results_or_values(input))
    _ods_context = _ods_get_default_loc_context(loc)
    results.append(output)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def input(self):
    _ods_variadic_group_length = len(self.operation.operands) - 1 + 1
    return self.operation.operands[0:0 + _ods_variadic_group_length]

  @builtins.property
  def output(self):
    return self.operation.results[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class RotateOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.rotate"

  _ODS_REGIONS = (0, True)

  def __init__(self, value, amount, dimension, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(value))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["amount"] = (amount if (
    issubclass(type(amount), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('SI32Attr')) else
      _ods_ir.AttrBuilder.get('SI32Attr')(amount, context=_ods_context))
    attributes["dimension"] = (dimension if (
    issubclass(type(dimension), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I32Attr')) else
      _ods_ir.AttrBuilder.get('I32Attr')(dimension, context=_ods_context))
    results.extend([operands[0].type] * 1)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def value(self):
    return self.operation.operands[0]

  @builtins.property
  def amount(self):
    return self.operation.attributes["amount"]

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

  @builtins.property
  def dimension(self):
    return self.operation.attributes["dimension"]

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

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

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class SemaphoreSignalOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.sem_signal"

  _ODS_REGIONS = (0, True)

  def __init__(self, semaphore, amount, *, device_id=None, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(semaphore))
    operands.append(_get_op_result_or_value(amount))
    if device_id is not None: operands.append(_get_op_result_or_value(device_id))
    _ods_context = _ods_get_default_loc_context(loc)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def semaphore(self):
    return self.operation.operands[0]

  @builtins.property
  def amount(self):
    return self.operation.operands[1]

  @builtins.property
  def device_id(self):
    return None if len(self.operation.operands) < 3 else self.operation.operands[2]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class SemaphoreWaitOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.sem_wait"

  _ODS_REGIONS = (0, True)

  def __init__(self, semaphore, amount, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(semaphore))
    operands.append(_get_op_result_or_value(amount))
    _ods_context = _ods_get_default_loc_context(loc)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def semaphore(self):
    return self.operation.operands[0]

  @builtins.property
  def amount(self):
    return self.operation.operands[1]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class StoreOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.store"

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

  _ODS_REGIONS = (0, True)

  def __init__(self, valueToStore, base, indices, sublane_mask, *, mask=None, sublane_stride=None, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(valueToStore))
    operands.append(_get_op_result_or_value(base))
    operands.append(_get_op_results_or_values(indices))
    operands.append(_get_op_result_or_value(mask) if mask is not None else None)
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["sublane_mask"] = (sublane_mask if (
    issubclass(type(sublane_mask), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('DenseBoolArrayAttr')) else
      _ods_ir.AttrBuilder.get('DenseBoolArrayAttr')(sublane_mask, context=_ods_context))
    if sublane_stride is not None: attributes["sublane_stride"] = (sublane_stride if (
        issubclass(type(sublane_stride), _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('I32Attr')) else
          _ods_ir.AttrBuilder.get('I32Attr')(sublane_stride, context=_ods_context))
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def valueToStore(self):
    operand_range = _ods_segmented_accessor(
         self.operation.operands,
         self.operation.attributes["operandSegmentSizes"], 0)
    return operand_range[0]

  @builtins.property
  def base(self):
    operand_range = _ods_segmented_accessor(
         self.operation.operands,
         self.operation.attributes["operandSegmentSizes"], 1)
    return operand_range[0]

  @builtins.property
  def indices(self):
    operand_range = _ods_segmented_accessor(
         self.operation.operands,
         self.operation.attributes["operandSegmentSizes"], 2)
    return operand_range

  @builtins.property
  def mask(self):
    operand_range = _ods_segmented_accessor(
         self.operation.operands,
         self.operation.attributes["operandSegmentSizes"], 3)
    return operand_range[0] if len(operand_range) > 0 else None

  @builtins.property
  def sublane_mask(self):
    return self.operation.attributes["sublane_mask"]

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

  @builtins.property
  def sublane_stride(self):
    if "sublane_stride" not in self.operation.attributes:
      return None
    return self.operation.attributes["sublane_stride"]

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

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

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class TraceOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.trace"

  _ODS_REGIONS = (1, True)

  def __init__(self, results_, message, level, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["message"] = (message if (
    issubclass(type(message), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('StrAttr')) else
      _ods_ir.AttrBuilder.get('StrAttr')(message, context=_ods_context))
    attributes["level"] = (level if (
    issubclass(type(level), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I32Attr')) else
      _ods_ir.AttrBuilder.get('I32Attr')(level, context=_ods_context))
    results.extend(results_)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def message(self):
    return self.operation.attributes["message"]

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

  @builtins.property
  def level(self):
    return self.operation.attributes["level"]

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

  @builtins.property
  def results_(self):
    _ods_variadic_group_length = len(self.operation.results) - 1 + 1
    return self.operation.results[0:0 + _ods_variadic_group_length]

  @builtins.property
  def region(self):
    return self.regions[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class TraceStartOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.trace_start"

  _ODS_REGIONS = (0, True)

  def __init__(self, message, level, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["message"] = (message if (
    issubclass(type(message), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('StrAttr')) else
      _ods_ir.AttrBuilder.get('StrAttr')(message, context=_ods_context))
    attributes["level"] = (level if (
    issubclass(type(level), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I32Attr')) else
      _ods_ir.AttrBuilder.get('I32Attr')(level, context=_ods_context))
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def message(self):
    return self.operation.attributes["message"]

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

  @builtins.property
  def level(self):
    return self.operation.attributes["level"]

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

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class TraceStopOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.trace_stop"

  _ODS_REGIONS = (0, True)

  def __init__(self, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class UnpackSubelementsOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.unpack_subelements"

  _ODS_REGIONS = (0, True)

  def __init__(self, output, source, index, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(source))
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["index"] = (index if (
    issubclass(type(index), _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I32Attr')) else
      _ods_ir.AttrBuilder.get('I32Attr')(index, context=_ods_context))
    results.append(output)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def source(self):
    return self.operation.operands[0]

  @builtins.property
  def index(self):
    return self.operation.attributes["index"]

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

  @builtins.property
  def output(self):
    return self.operation.results[0]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class UnrollVectorsOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.unroll_vectors"

  _ODS_REGIONS = (0, True)

  def __init__(self, output, input, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(input))
    _ods_context = _ods_get_default_loc_context(loc)
    results.extend(output)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

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

  @builtins.property
  def output(self):
    _ods_variadic_group_length = len(self.operation.results) - 1 + 1
    return self.operation.results[0:0 + _ods_variadic_group_length]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class WaitDMAOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.wait_dma"

  _ODS_REGIONS = (0, True)

  def __init__(self, semaphore, ref, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.append(_get_op_result_or_value(semaphore))
    operands.append(_get_op_result_or_value(ref))
    _ods_context = _ods_get_default_loc_context(loc)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def semaphore(self):
    return self.operation.operands[0]

  @builtins.property
  def ref(self):
    return self.operation.operands[1]

@_ods_cext.register_operation(_Dialect)
@_ods_extend_opview_class(_ods_ext_module)
class YieldOp(_ods_ir.OpView):
  OPERATION_NAME = "tpu.yield"

  _ODS_REGIONS = (0, True)

  def __init__(self, results_, *, loc=None, ip=None):
    operands = []
    results = []
    attributes = {}
    regions = None
    operands.extend(_get_op_results_or_values(results_))
    _ods_context = _ods_get_default_loc_context(loc)
    _ods_successors = None
    super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))

  @builtins.property
  def results_(self):
    _ods_variadic_group_length = len(self.operation.operands) - 1 + 1
    return self.operation.operands[0:0 + _ods_variadic_group_length]
