# 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_result_or_op_results as _get_op_result_or_op_results, get_op_result_or_value as _get_op_result_or_value, get_op_results_or_values as _get_op_results_or_values, segmented_accessor as _ods_segmented_accessor, ) _ods_ir = _ods_cext.ir import builtins from typing import Sequence as _Sequence, Union as _Union @_ods_cext.register_dialect class _Dialect(_ods_ir.Dialect): DIALECT_NAMESPACE = "mhlo" @_ods_cext.register_operation(_Dialect) class AbsOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.abs" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def abs(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(AbsOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class AddDependencyOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.add_dependency" _ODS_REGIONS = (0, True) def __init__(self, operand, token, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(token)) _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 operand(self): return self.operation.operands[0] @builtins.property def token(self): return self.operation.operands[1] @builtins.property def output(self): return self.operation.results[0] def add_dependency(operand, token, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(AddDependencyOp(operand=operand, token=token, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class AddOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.add" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def add(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(AddOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class AfterAllOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.after_all" _ODS_REGIONS = (0, True) def __init__(self, inputs, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(inputs)) _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 inputs(self): _ods_variadic_group_length = len(self.operation.operands) - 1 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @builtins.property def result(self): return self.operation.results[0] def after_all(inputs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(AfterAllOp(inputs=inputs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class AllGatherOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.all_gather" _ODS_REGIONS = (0, True) def __init__(self, result, operands_, all_gather_dim, replica_groups, *, channel_handle=None, use_global_device_ids=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(operands_)) _ods_context = _ods_get_default_loc_context(loc) attributes["all_gather_dim"] = (all_gather_dim if ( isinstance(all_gather_dim, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(all_gather_dim, context=_ods_context)) attributes["replica_groups"] = (replica_groups if ( isinstance(replica_groups, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(replica_groups, context=_ods_context)) if channel_handle is not None: attributes["channel_handle"] = (channel_handle if ( isinstance(channel_handle, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ChannelHandle')) else _ods_ir.AttrBuilder.get('MHLO_ChannelHandle')(channel_handle, context=_ods_context)) if bool(use_global_device_ids): attributes["use_global_device_ids"] = _ods_ir.UnitAttr.get( _ods_get_default_loc_context(loc)) results.extend(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 operands_(self): _ods_variadic_group_length = len(self.operation.operands) - 1 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @builtins.property def all_gather_dim(self): return self.operation.attributes["all_gather_dim"] @all_gather_dim.setter def all_gather_dim(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["all_gather_dim"] = value @builtins.property def replica_groups(self): return self.operation.attributes["replica_groups"] @replica_groups.setter def replica_groups(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["replica_groups"] = value @builtins.property def channel_handle(self): if "channel_handle" not in self.operation.attributes: return None return self.operation.attributes["channel_handle"] @channel_handle.setter def channel_handle(self, value): if value is not None: self.operation.attributes["channel_handle"] = value elif "channel_handle" in self.operation.attributes: del self.operation.attributes["channel_handle"] @channel_handle.deleter def channel_handle(self): del self.operation.attributes["channel_handle"] @builtins.property def use_global_device_ids(self): return "use_global_device_ids" in self.operation.attributes @use_global_device_ids.setter def use_global_device_ids(self, value): if bool(value): self.operation.attributes["use_global_device_ids"] = _ods_ir.UnitAttr.get() elif "use_global_device_ids" in self.operation.attributes: del self.operation.attributes["use_global_device_ids"] @use_global_device_ids.deleter def use_global_device_ids(self): del self.operation.attributes["use_global_device_ids"] def all_gather(result, operands_, all_gather_dim, replica_groups, *, channel_handle=None, use_global_device_ids=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(AllGatherOp(result=result, operands_=operands_, all_gather_dim=all_gather_dim, replica_groups=replica_groups, channel_handle=channel_handle, use_global_device_ids=use_global_device_ids, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class AllReduceOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.all_reduce" _ODS_REGIONS = (1, True) def __init__(self, result, operands_, replica_groups, *, channel_handle=None, use_global_device_ids=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(operands_)) _ods_context = _ods_get_default_loc_context(loc) attributes["replica_groups"] = (replica_groups if ( isinstance(replica_groups, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(replica_groups, context=_ods_context)) if channel_handle is not None: attributes["channel_handle"] = (channel_handle if ( isinstance(channel_handle, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ChannelHandle')) else _ods_ir.AttrBuilder.get('MHLO_ChannelHandle')(channel_handle, context=_ods_context)) if bool(use_global_device_ids): attributes["use_global_device_ids"] = _ods_ir.UnitAttr.get( _ods_get_default_loc_context(loc)) results.extend(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 operands_(self): _ods_variadic_group_length = len(self.operation.operands) - 1 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @builtins.property def replica_groups(self): return self.operation.attributes["replica_groups"] @replica_groups.setter def replica_groups(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["replica_groups"] = value @builtins.property def channel_handle(self): if "channel_handle" not in self.operation.attributes: return None return self.operation.attributes["channel_handle"] @channel_handle.setter def channel_handle(self, value): if value is not None: self.operation.attributes["channel_handle"] = value elif "channel_handle" in self.operation.attributes: del self.operation.attributes["channel_handle"] @channel_handle.deleter def channel_handle(self): del self.operation.attributes["channel_handle"] @builtins.property def use_global_device_ids(self): return "use_global_device_ids" in self.operation.attributes @use_global_device_ids.setter def use_global_device_ids(self, value): if bool(value): self.operation.attributes["use_global_device_ids"] = _ods_ir.UnitAttr.get() elif "use_global_device_ids" in self.operation.attributes: del self.operation.attributes["use_global_device_ids"] @use_global_device_ids.deleter def use_global_device_ids(self): del self.operation.attributes["use_global_device_ids"] @builtins.property def computation(self): return self.regions[0] def all_reduce(result, operands_, replica_groups, *, channel_handle=None, use_global_device_ids=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(AllReduceOp(result=result, operands_=operands_, replica_groups=replica_groups, channel_handle=channel_handle, use_global_device_ids=use_global_device_ids, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class AllToAllOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.all_to_all" _ODS_REGIONS = (0, True) def __init__(self, operand, replica_groups, *, split_dimension=None, concat_dimension=None, split_count=None, channel_handle=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(operand)) _ods_context = _ods_get_default_loc_context(loc) if split_dimension is not None: attributes["split_dimension"] = (split_dimension if ( isinstance(split_dimension, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(split_dimension, context=_ods_context)) if concat_dimension is not None: attributes["concat_dimension"] = (concat_dimension if ( isinstance(concat_dimension, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(concat_dimension, context=_ods_context)) if split_count is not None: attributes["split_count"] = (split_count if ( isinstance(split_count, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(split_count, context=_ods_context)) attributes["replica_groups"] = (replica_groups if ( isinstance(replica_groups, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(replica_groups, context=_ods_context)) if channel_handle is not None: attributes["channel_handle"] = (channel_handle if ( isinstance(channel_handle, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ChannelHandle')) else _ods_ir.AttrBuilder.get('MHLO_ChannelHandle')(channel_handle, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): _ods_variadic_group_length = len(self.operation.operands) - 1 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @builtins.property def split_dimension(self): if "split_dimension" not in self.operation.attributes: return None return self.operation.attributes["split_dimension"] @split_dimension.setter def split_dimension(self, value): if value is not None: self.operation.attributes["split_dimension"] = value elif "split_dimension" in self.operation.attributes: del self.operation.attributes["split_dimension"] @split_dimension.deleter def split_dimension(self): del self.operation.attributes["split_dimension"] @builtins.property def concat_dimension(self): if "concat_dimension" not in self.operation.attributes: return None return self.operation.attributes["concat_dimension"] @concat_dimension.setter def concat_dimension(self, value): if value is not None: self.operation.attributes["concat_dimension"] = value elif "concat_dimension" in self.operation.attributes: del self.operation.attributes["concat_dimension"] @concat_dimension.deleter def concat_dimension(self): del self.operation.attributes["concat_dimension"] @builtins.property def split_count(self): if "split_count" not in self.operation.attributes: return None return self.operation.attributes["split_count"] @split_count.setter def split_count(self, value): if value is not None: self.operation.attributes["split_count"] = value elif "split_count" in self.operation.attributes: del self.operation.attributes["split_count"] @split_count.deleter def split_count(self): del self.operation.attributes["split_count"] @builtins.property def replica_groups(self): return self.operation.attributes["replica_groups"] @replica_groups.setter def replica_groups(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["replica_groups"] = value @builtins.property def channel_handle(self): if "channel_handle" not in self.operation.attributes: return None return self.operation.attributes["channel_handle"] @channel_handle.setter def channel_handle(self, value): if value is not None: self.operation.attributes["channel_handle"] = value elif "channel_handle" in self.operation.attributes: del self.operation.attributes["channel_handle"] @channel_handle.deleter def channel_handle(self): del self.operation.attributes["channel_handle"] def all_to_all(operand, replica_groups, *, split_dimension=None, concat_dimension=None, split_count=None, channel_handle=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(AllToAllOp(operand=operand, replica_groups=replica_groups, split_dimension=split_dimension, concat_dimension=concat_dimension, split_count=split_count, channel_handle=channel_handle, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class AndOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.and" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def and_(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(AndOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class AsyncDoneOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.async_done" _ODS_REGIONS = (0, True) def __init__(self, bundle, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(bundle)) _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 bundle(self): return self.operation.operands[0] def async_done(bundle, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(AsyncDoneOp(bundle=bundle, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class AsyncStartOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.async_start" _ODS_REGIONS = (0, True) def __init__(self, result, inputs, called_computation, execution_thread, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(inputs)) _ods_context = _ods_get_default_loc_context(loc) attributes["called_computation"] = (called_computation if ( isinstance(called_computation, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(called_computation, context=_ods_context)) attributes["execution_thread"] = (execution_thread if ( isinstance(execution_thread, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('StrAttr')) else _ods_ir.AttrBuilder.get('StrAttr')(execution_thread, 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 inputs(self): _ods_variadic_group_length = len(self.operation.operands) - 1 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @builtins.property def called_computation(self): return self.operation.attributes["called_computation"] @called_computation.setter def called_computation(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["called_computation"] = value @builtins.property def execution_thread(self): return self.operation.attributes["execution_thread"] @execution_thread.setter def execution_thread(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["execution_thread"] = value def async_start(result, inputs, called_computation, execution_thread, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(AsyncStartOp(result=result, inputs=inputs, called_computation=called_computation, execution_thread=execution_thread, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class AsyncUpdateOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.async_update" _ODS_REGIONS = (0, True) def __init__(self, bundle, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(bundle)) _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 bundle(self): return self.operation.operands[0] def async_update(bundle, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(AsyncUpdateOp(bundle=bundle, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class Atan2Op(_ods_ir.OpView): OPERATION_NAME = "mhlo.atan2" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def atan2(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(Atan2Op(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class BatchNormGradOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.batch_norm_grad" _ODS_REGIONS = (0, True) def __init__(self, operand, scale, mean, variance, grad_output, epsilon, feature_index, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(scale)) operands.append(_get_op_result_or_value(mean)) operands.append(_get_op_result_or_value(variance)) operands.append(_get_op_result_or_value(grad_output)) _ods_context = _ods_get_default_loc_context(loc) attributes["epsilon"] = (epsilon if ( isinstance(epsilon, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('F32Attr')) else _ods_ir.AttrBuilder.get('F32Attr')(epsilon, context=_ods_context)) attributes["feature_index"] = (feature_index if ( isinstance(feature_index, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(feature_index, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def scale(self): return self.operation.operands[1] @builtins.property def mean(self): return self.operation.operands[2] @builtins.property def variance(self): return self.operation.operands[3] @builtins.property def grad_output(self): return self.operation.operands[4] @builtins.property def epsilon(self): return self.operation.attributes["epsilon"] @epsilon.setter def epsilon(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["epsilon"] = value @builtins.property def feature_index(self): return self.operation.attributes["feature_index"] @feature_index.setter def feature_index(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["feature_index"] = value @builtins.property def grad_operand(self): return self.operation.results[0] @builtins.property def grad_scale(self): return self.operation.results[1] @builtins.property def grad_offset(self): return self.operation.results[2] def batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index, *, loc=None, ip=None) -> _Sequence[_ods_ir.Value]: return _get_op_result_or_op_results(BatchNormGradOp(operand=operand, scale=scale, mean=mean, variance=variance, grad_output=grad_output, epsilon=epsilon, feature_index=feature_index, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class BatchNormInferenceOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.batch_norm_inference" _ODS_REGIONS = (0, True) def __init__(self, operand, scale, offset, mean, variance, epsilon, feature_index, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(scale)) operands.append(_get_op_result_or_value(offset)) operands.append(_get_op_result_or_value(mean)) operands.append(_get_op_result_or_value(variance)) _ods_context = _ods_get_default_loc_context(loc) attributes["epsilon"] = (epsilon if ( isinstance(epsilon, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('F32Attr')) else _ods_ir.AttrBuilder.get('F32Attr')(epsilon, context=_ods_context)) attributes["feature_index"] = (feature_index if ( isinstance(feature_index, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(feature_index, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def scale(self): return self.operation.operands[1] @builtins.property def offset(self): return self.operation.operands[2] @builtins.property def mean(self): return self.operation.operands[3] @builtins.property def variance(self): return self.operation.operands[4] @builtins.property def epsilon(self): return self.operation.attributes["epsilon"] @epsilon.setter def epsilon(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["epsilon"] = value @builtins.property def feature_index(self): return self.operation.attributes["feature_index"] @feature_index.setter def feature_index(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["feature_index"] = value @builtins.property def result(self): return self.operation.results[0] def batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(BatchNormInferenceOp(operand=operand, scale=scale, offset=offset, mean=mean, variance=variance, epsilon=epsilon, feature_index=feature_index, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class BatchNormTrainingOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.batch_norm_training" _ODS_REGIONS = (0, True) def __init__(self, operand, scale, offset, epsilon, feature_index, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(scale)) operands.append(_get_op_result_or_value(offset)) _ods_context = _ods_get_default_loc_context(loc) attributes["epsilon"] = (epsilon if ( isinstance(epsilon, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('F32Attr')) else _ods_ir.AttrBuilder.get('F32Attr')(epsilon, context=_ods_context)) attributes["feature_index"] = (feature_index if ( isinstance(feature_index, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(feature_index, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def scale(self): return self.operation.operands[1] @builtins.property def offset(self): return self.operation.operands[2] @builtins.property def epsilon(self): return self.operation.attributes["epsilon"] @epsilon.setter def epsilon(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["epsilon"] = value @builtins.property def feature_index(self): return self.operation.attributes["feature_index"] @feature_index.setter def feature_index(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["feature_index"] = value @builtins.property def output(self): return self.operation.results[0] @builtins.property def batch_mean(self): return self.operation.results[1] @builtins.property def batch_var(self): return self.operation.results[2] def batch_norm_training(operand, scale, offset, epsilon, feature_index, *, loc=None, ip=None) -> _Sequence[_ods_ir.Value]: return _get_op_result_or_op_results(BatchNormTrainingOp(operand=operand, scale=scale, offset=offset, epsilon=epsilon, feature_index=feature_index, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class BitcastConvertOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.bitcast_convert" _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] def bitcast_convert(result, operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(BitcastConvertOp(result=result, operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class BitcastOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.bitcast" _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] def bitcast(result, operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(BitcastOp(result=result, operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class BroadcastInDimOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.broadcast_in_dim" _ODS_REGIONS = (0, True) def __init__(self, result, operand, broadcast_dimensions, *, 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) attributes["broadcast_dimensions"] = (broadcast_dimensions if ( isinstance(broadcast_dimensions, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(broadcast_dimensions, 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 operand(self): return self.operation.operands[0] @builtins.property def broadcast_dimensions(self): return self.operation.attributes["broadcast_dimensions"] @broadcast_dimensions.setter def broadcast_dimensions(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["broadcast_dimensions"] = value def broadcast_in_dim(result, operand, broadcast_dimensions, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(BroadcastInDimOp(result=result, operand=operand, broadcast_dimensions=broadcast_dimensions, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class BroadcastOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.broadcast" _ODS_REGIONS = (0, True) def __init__(self, operand, broadcast_sizes, *, 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) attributes["broadcast_sizes"] = (broadcast_sizes if ( isinstance(broadcast_sizes, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(broadcast_sizes, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def broadcast_sizes(self): return self.operation.attributes["broadcast_sizes"] @broadcast_sizes.setter def broadcast_sizes(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["broadcast_sizes"] = value def broadcast(operand, broadcast_sizes, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(BroadcastOp(operand=operand, broadcast_sizes=broadcast_sizes, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class CaseOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.case" _ODS_REGIONS = (0, False) def __init__(self, result, index, num_branches, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(index)) _ods_context = _ods_get_default_loc_context(loc) results.extend(result) _ods_successors = None regions = 0 + num_branches super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def index(self): return self.operation.operands[0] @builtins.property def branches(self): return self.regions[0:] def case(result, index, num_branches, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(CaseOp(result=result, index=index, num_branches=num_branches, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class CbrtOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.cbrt" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def cbrt(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(CbrtOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class CeilOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.ceil" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def ceil(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(CeilOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class CholeskyOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.cholesky" _ODS_REGIONS = (0, True) def __init__(self, a, *, lower=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(a)) _ods_context = _ods_get_default_loc_context(loc) if lower is not None: attributes["lower"] = (lower if ( isinstance(lower, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('BoolAttr')) else _ods_ir.AttrBuilder.get('BoolAttr')(lower, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def a(self): return self.operation.operands[0] @builtins.property def lower(self): if "lower" not in self.operation.attributes: return None return self.operation.attributes["lower"] @lower.setter def lower(self, value): if value is not None: self.operation.attributes["lower"] = value elif "lower" in self.operation.attributes: del self.operation.attributes["lower"] @lower.deleter def lower(self): del self.operation.attributes["lower"] def cholesky(a, *, lower=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(CholeskyOp(a=a, lower=lower, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ClampOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.clamp" _ODS_REGIONS = (0, True) def __init__(self, min, operand, max, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(min)) operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(max)) _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 min(self): return self.operation.operands[0] @builtins.property def operand(self): return self.operation.operands[1] @builtins.property def max(self): return self.operation.operands[2] @builtins.property def result(self): return self.operation.results[0] def clamp(min, operand, max, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ClampOp(min=min, operand=operand, max=max, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ClzOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.count_leading_zeros" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def count_leading_zeros(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ClzOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class CollectiveBroadcastOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.collective_broadcast" _ODS_REGIONS = (0, True) def __init__(self, operand, replica_groups, *, channel_handle=None, 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) attributes["replica_groups"] = (replica_groups if ( isinstance(replica_groups, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(replica_groups, context=_ods_context)) if channel_handle is not None: attributes["channel_handle"] = (channel_handle if ( isinstance(channel_handle, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ChannelHandle')) else _ods_ir.AttrBuilder.get('MHLO_ChannelHandle')(channel_handle, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def replica_groups(self): return self.operation.attributes["replica_groups"] @replica_groups.setter def replica_groups(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["replica_groups"] = value @builtins.property def channel_handle(self): if "channel_handle" not in self.operation.attributes: return None return self.operation.attributes["channel_handle"] @channel_handle.setter def channel_handle(self, value): if value is not None: self.operation.attributes["channel_handle"] = value elif "channel_handle" in self.operation.attributes: del self.operation.attributes["channel_handle"] @channel_handle.deleter def channel_handle(self): del self.operation.attributes["channel_handle"] def collective_broadcast(operand, replica_groups, *, channel_handle=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(CollectiveBroadcastOp(operand=operand, replica_groups=replica_groups, channel_handle=channel_handle, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class CollectivePermuteOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.collective_permute" _ODS_REGIONS = (0, True) def __init__(self, operand, source_target_pairs, *, channel_handle=None, 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) attributes["source_target_pairs"] = (source_target_pairs if ( isinstance(source_target_pairs, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(source_target_pairs, context=_ods_context)) if channel_handle is not None: attributes["channel_handle"] = (channel_handle if ( isinstance(channel_handle, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ChannelHandle')) else _ods_ir.AttrBuilder.get('MHLO_ChannelHandle')(channel_handle, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def source_target_pairs(self): return self.operation.attributes["source_target_pairs"] @source_target_pairs.setter def source_target_pairs(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["source_target_pairs"] = value @builtins.property def channel_handle(self): if "channel_handle" not in self.operation.attributes: return None return self.operation.attributes["channel_handle"] @channel_handle.setter def channel_handle(self, value): if value is not None: self.operation.attributes["channel_handle"] = value elif "channel_handle" in self.operation.attributes: del self.operation.attributes["channel_handle"] @channel_handle.deleter def channel_handle(self): del self.operation.attributes["channel_handle"] def collective_permute(operand, source_target_pairs, *, channel_handle=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(CollectivePermuteOp(operand=operand, source_target_pairs=source_target_pairs, channel_handle=channel_handle, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class CompareOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.compare" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, comparison_direction, *, compare_type=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)) _ods_context = _ods_get_default_loc_context(loc) attributes["comparison_direction"] = (comparison_direction if ( isinstance(comparison_direction, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ComparisonDirectionAttr')) else _ods_ir.AttrBuilder.get('MHLO_ComparisonDirectionAttr')(comparison_direction, context=_ods_context)) if compare_type is not None: attributes["compare_type"] = (compare_type if ( isinstance(compare_type, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ComparisonTypeAttr')) else _ods_ir.AttrBuilder.get('MHLO_ComparisonTypeAttr')(compare_type, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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 comparison_direction(self): return self.operation.attributes["comparison_direction"] @comparison_direction.setter def comparison_direction(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["comparison_direction"] = value @builtins.property def compare_type(self): if "compare_type" not in self.operation.attributes: return None return self.operation.attributes["compare_type"] @compare_type.setter def compare_type(self, value): if value is not None: self.operation.attributes["compare_type"] = value elif "compare_type" in self.operation.attributes: del self.operation.attributes["compare_type"] @compare_type.deleter def compare_type(self): del self.operation.attributes["compare_type"] def compare(lhs, rhs, comparison_direction, *, compare_type=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(CompareOp(lhs=lhs, rhs=rhs, comparison_direction=comparison_direction, compare_type=compare_type, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ComplexOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.complex" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def complex(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ComplexOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class CompositeOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.composite" _ODS_REGIONS = (0, True) def __init__(self, result, inputs, name, decomposition, *, composite_attributes=None, version=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(inputs)) _ods_context = _ods_get_default_loc_context(loc) attributes["name"] = (name if ( isinstance(name, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('StrAttr')) else _ods_ir.AttrBuilder.get('StrAttr')(name, context=_ods_context)) if composite_attributes is not None: attributes["composite_attributes"] = (composite_attributes if ( isinstance(composite_attributes, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('DictionaryAttr')) else _ods_ir.AttrBuilder.get('DictionaryAttr')(composite_attributes, context=_ods_context)) attributes["decomposition"] = (decomposition if ( isinstance(decomposition, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(decomposition, context=_ods_context)) if version is not None: attributes["version"] = (version if ( isinstance(version, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I32Attr')) else _ods_ir.AttrBuilder.get('I32Attr')(version, context=_ods_context)) results.extend(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 inputs(self): _ods_variadic_group_length = len(self.operation.operands) - 1 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @builtins.property def name(self): return self.operation.attributes["name"] @name.setter def name(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["name"] = value @builtins.property def composite_attributes(self): if "composite_attributes" not in self.operation.attributes: return None return self.operation.attributes["composite_attributes"] @composite_attributes.setter def composite_attributes(self, value): if value is not None: self.operation.attributes["composite_attributes"] = value elif "composite_attributes" in self.operation.attributes: del self.operation.attributes["composite_attributes"] @composite_attributes.deleter def composite_attributes(self): del self.operation.attributes["composite_attributes"] @builtins.property def decomposition(self): return self.operation.attributes["decomposition"] @decomposition.setter def decomposition(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["decomposition"] = value @builtins.property def version(self): if "version" not in self.operation.attributes: return None return self.operation.attributes["version"] @version.setter def version(self, value): if value is not None: self.operation.attributes["version"] = value elif "version" in self.operation.attributes: del self.operation.attributes["version"] @version.deleter def version(self): del self.operation.attributes["version"] def composite(result, inputs, name, decomposition, *, composite_attributes=None, version=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(CompositeOp(result=result, inputs=inputs, name=name, decomposition=decomposition, composite_attributes=composite_attributes, version=version, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ConcatenateOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.concatenate" _ODS_REGIONS = (0, True) def __init__(self, val, dimension, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(val)) _ods_context = _ods_get_default_loc_context(loc) attributes["dimension"] = (dimension if ( isinstance(dimension, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(dimension, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def val(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 def concatenate(val, dimension, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ConcatenateOp(val=val, dimension=dimension, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ConstantOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.constant" _ODS_REGIONS = (0, True) def __init__(self, value, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None _ods_context = _ods_get_default_loc_context(loc) attributes["value"] = (value if ( isinstance(value, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('ElementsAttr')) else _ods_ir.AttrBuilder.get('ElementsAttr')(value, context=_ods_context)) _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.attributes["value"] @value.setter def value(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["value"] = value @builtins.property def output(self): return self.operation.results[0] def constant(value, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ConstantOp(value=value, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ConvertOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.convert" _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] def convert(result, operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ConvertOp(result=result, operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ConvolutionOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.convolution" _ODS_REGIONS = (0, True) def __init__(self, result, lhs, rhs, dimension_numbers, feature_group_count, batch_group_count, *, window_strides=None, padding=None, lhs_dilation=None, rhs_dilation=None, window_reversal=None, precision_config=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)) _ods_context = _ods_get_default_loc_context(loc) if window_strides is not None: attributes["window_strides"] = (window_strides if ( isinstance(window_strides, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(window_strides, context=_ods_context)) if padding is not None: attributes["padding"] = (padding if ( isinstance(padding, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(padding, context=_ods_context)) if lhs_dilation is not None: attributes["lhs_dilation"] = (lhs_dilation if ( isinstance(lhs_dilation, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(lhs_dilation, context=_ods_context)) if rhs_dilation is not None: attributes["rhs_dilation"] = (rhs_dilation if ( isinstance(rhs_dilation, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(rhs_dilation, context=_ods_context)) if window_reversal is not None: attributes["window_reversal"] = (window_reversal if ( isinstance(window_reversal, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_BoolElementsAttr')) else _ods_ir.AttrBuilder.get('MHLO_BoolElementsAttr')(window_reversal, context=_ods_context)) attributes["dimension_numbers"] = (dimension_numbers if ( isinstance(dimension_numbers, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ConvDimensionNumbers')) else _ods_ir.AttrBuilder.get('MHLO_ConvDimensionNumbers')(dimension_numbers, context=_ods_context)) attributes["feature_group_count"] = (feature_group_count if ( isinstance(feature_group_count, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(feature_group_count, context=_ods_context)) attributes["batch_group_count"] = (batch_group_count if ( isinstance(batch_group_count, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(batch_group_count, context=_ods_context)) if precision_config is not None: attributes["precision_config"] = (precision_config if ( isinstance(precision_config, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_PrecisionConfigAttr')) else _ods_ir.AttrBuilder.get('MHLO_PrecisionConfigAttr')(precision_config, 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 window_strides(self): if "window_strides" not in self.operation.attributes: return None return self.operation.attributes["window_strides"] @window_strides.setter def window_strides(self, value): if value is not None: self.operation.attributes["window_strides"] = value elif "window_strides" in self.operation.attributes: del self.operation.attributes["window_strides"] @window_strides.deleter def window_strides(self): del self.operation.attributes["window_strides"] @builtins.property def padding(self): if "padding" not in self.operation.attributes: return None return self.operation.attributes["padding"] @padding.setter def padding(self, value): if value is not None: self.operation.attributes["padding"] = value elif "padding" in self.operation.attributes: del self.operation.attributes["padding"] @padding.deleter def padding(self): del self.operation.attributes["padding"] @builtins.property def lhs_dilation(self): if "lhs_dilation" not in self.operation.attributes: return None return self.operation.attributes["lhs_dilation"] @lhs_dilation.setter def lhs_dilation(self, value): if value is not None: self.operation.attributes["lhs_dilation"] = value elif "lhs_dilation" in self.operation.attributes: del self.operation.attributes["lhs_dilation"] @lhs_dilation.deleter def lhs_dilation(self): del self.operation.attributes["lhs_dilation"] @builtins.property def rhs_dilation(self): if "rhs_dilation" not in self.operation.attributes: return None return self.operation.attributes["rhs_dilation"] @rhs_dilation.setter def rhs_dilation(self, value): if value is not None: self.operation.attributes["rhs_dilation"] = value elif "rhs_dilation" in self.operation.attributes: del self.operation.attributes["rhs_dilation"] @rhs_dilation.deleter def rhs_dilation(self): del self.operation.attributes["rhs_dilation"] @builtins.property def window_reversal(self): if "window_reversal" not in self.operation.attributes: return None return self.operation.attributes["window_reversal"] @window_reversal.setter def window_reversal(self, value): if value is not None: self.operation.attributes["window_reversal"] = value elif "window_reversal" in self.operation.attributes: del self.operation.attributes["window_reversal"] @window_reversal.deleter def window_reversal(self): del self.operation.attributes["window_reversal"] @builtins.property def dimension_numbers(self): return self.operation.attributes["dimension_numbers"] @dimension_numbers.setter def dimension_numbers(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["dimension_numbers"] = value @builtins.property def feature_group_count(self): return self.operation.attributes["feature_group_count"] @feature_group_count.setter def feature_group_count(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["feature_group_count"] = value @builtins.property def batch_group_count(self): return self.operation.attributes["batch_group_count"] @batch_group_count.setter def batch_group_count(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["batch_group_count"] = value @builtins.property def precision_config(self): if "precision_config" not in self.operation.attributes: return None return self.operation.attributes["precision_config"] @precision_config.setter def precision_config(self, value): if value is not None: self.operation.attributes["precision_config"] = value elif "precision_config" in self.operation.attributes: del self.operation.attributes["precision_config"] @precision_config.deleter def precision_config(self): del self.operation.attributes["precision_config"] def convolution(result, lhs, rhs, dimension_numbers, feature_group_count, batch_group_count, *, window_strides=None, padding=None, lhs_dilation=None, rhs_dilation=None, window_reversal=None, precision_config=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ConvolutionOp(result=result, lhs=lhs, rhs=rhs, dimension_numbers=dimension_numbers, feature_group_count=feature_group_count, batch_group_count=batch_group_count, window_strides=window_strides, padding=padding, lhs_dilation=lhs_dilation, rhs_dilation=rhs_dilation, window_reversal=window_reversal, precision_config=precision_config, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class CopyOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.copy" _ODS_REGIONS = (0, True) def __init__(self, operand, *, cross_program_prefetch_index=None, 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) if cross_program_prefetch_index is not None: attributes["cross_program_prefetch_index"] = (cross_program_prefetch_index if ( isinstance(cross_program_prefetch_index, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I32Attr')) else _ods_ir.AttrBuilder.get('I32Attr')(cross_program_prefetch_index, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def cross_program_prefetch_index(self): if "cross_program_prefetch_index" not in self.operation.attributes: return None return self.operation.attributes["cross_program_prefetch_index"] @cross_program_prefetch_index.setter def cross_program_prefetch_index(self, value): if value is not None: self.operation.attributes["cross_program_prefetch_index"] = value elif "cross_program_prefetch_index" in self.operation.attributes: del self.operation.attributes["cross_program_prefetch_index"] @cross_program_prefetch_index.deleter def cross_program_prefetch_index(self): del self.operation.attributes["cross_program_prefetch_index"] @builtins.property def result(self): return self.operation.results[0] def copy(operand, *, cross_program_prefetch_index=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(CopyOp(operand=operand, cross_program_prefetch_index=cross_program_prefetch_index, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class CosineOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.cosine" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def cosine(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(CosineOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class CreateTokenOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.create_token" _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 output(self): return self.operation.results[0] def create_token(*, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(CreateTokenOp(loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class CrossReplicaSumOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.cross-replica-sum" _ODS_REGIONS = (0, True) def __init__(self, operand, replica_groups, *, 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) attributes["replica_groups"] = (replica_groups if ( isinstance(replica_groups, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(replica_groups, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def replica_groups(self): return self.operation.attributes["replica_groups"] @replica_groups.setter def replica_groups(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["replica_groups"] = value def cross_replica_sum(operand, replica_groups, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(CrossReplicaSumOp(operand=operand, replica_groups=replica_groups, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class CustomCallOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.custom_call" _ODS_REGIONS = (0, True) def __init__(self, result, inputs, call_target_name, *, has_side_effect=None, backend_config=None, api_version=None, called_computations=None, custom_call_schedule=None, operand_layouts=None, result_layouts=None, output_operand_aliases=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(inputs)) _ods_context = _ods_get_default_loc_context(loc) attributes["call_target_name"] = (call_target_name if ( isinstance(call_target_name, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('StrAttr')) else _ods_ir.AttrBuilder.get('StrAttr')(call_target_name, context=_ods_context)) if has_side_effect is not None: attributes["has_side_effect"] = (has_side_effect if ( isinstance(has_side_effect, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('BoolAttr')) else _ods_ir.AttrBuilder.get('BoolAttr')(has_side_effect, context=_ods_context)) if backend_config is not None: attributes["backend_config"] = (backend_config if ( isinstance(backend_config, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('anonymous_807')) else _ods_ir.AttrBuilder.get('anonymous_807')(backend_config, context=_ods_context)) if api_version is not None: attributes["api_version"] = (api_version if ( isinstance(api_version, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_CustomCallApiVersionAttr')) else _ods_ir.AttrBuilder.get('MHLO_CustomCallApiVersionAttr')(api_version, context=_ods_context)) if called_computations is not None: attributes["called_computations"] = (called_computations if ( isinstance(called_computations, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_FlatSymbolRefArrayAttr')) else _ods_ir.AttrBuilder.get('MHLO_FlatSymbolRefArrayAttr')(called_computations, context=_ods_context)) if custom_call_schedule is not None: attributes["custom_call_schedule"] = (custom_call_schedule if ( isinstance(custom_call_schedule, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_CustomCallScheduleAttr')) else _ods_ir.AttrBuilder.get('MHLO_CustomCallScheduleAttr')(custom_call_schedule, context=_ods_context)) if operand_layouts is not None: attributes["operand_layouts"] = (operand_layouts if ( isinstance(operand_layouts, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ArrayOfLayoutAttr')) else _ods_ir.AttrBuilder.get('MHLO_ArrayOfLayoutAttr')(operand_layouts, context=_ods_context)) if result_layouts is not None: attributes["result_layouts"] = (result_layouts if ( isinstance(result_layouts, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ArrayOfLayoutAttr')) else _ods_ir.AttrBuilder.get('MHLO_ArrayOfLayoutAttr')(result_layouts, context=_ods_context)) if output_operand_aliases is not None: attributes["output_operand_aliases"] = (output_operand_aliases if ( isinstance(output_operand_aliases, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('anonymous_814')) else _ods_ir.AttrBuilder.get('anonymous_814')(output_operand_aliases, context=_ods_context)) results.extend(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 inputs(self): _ods_variadic_group_length = len(self.operation.operands) - 1 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @builtins.property def call_target_name(self): return self.operation.attributes["call_target_name"] @call_target_name.setter def call_target_name(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["call_target_name"] = value @builtins.property def has_side_effect(self): if "has_side_effect" not in self.operation.attributes: return None return self.operation.attributes["has_side_effect"] @has_side_effect.setter def has_side_effect(self, value): if value is not None: self.operation.attributes["has_side_effect"] = value elif "has_side_effect" in self.operation.attributes: del self.operation.attributes["has_side_effect"] @has_side_effect.deleter def has_side_effect(self): del self.operation.attributes["has_side_effect"] @builtins.property def backend_config(self): if "backend_config" not in self.operation.attributes: return None return self.operation.attributes["backend_config"] @backend_config.setter def backend_config(self, value): if value is not None: self.operation.attributes["backend_config"] = value elif "backend_config" in self.operation.attributes: del self.operation.attributes["backend_config"] @backend_config.deleter def backend_config(self): del self.operation.attributes["backend_config"] @builtins.property def api_version(self): if "api_version" not in self.operation.attributes: return None return self.operation.attributes["api_version"] @api_version.setter def api_version(self, value): if value is not None: self.operation.attributes["api_version"] = value elif "api_version" in self.operation.attributes: del self.operation.attributes["api_version"] @api_version.deleter def api_version(self): del self.operation.attributes["api_version"] @builtins.property def called_computations(self): if "called_computations" not in self.operation.attributes: return None return self.operation.attributes["called_computations"] @called_computations.setter def called_computations(self, value): if value is not None: self.operation.attributes["called_computations"] = value elif "called_computations" in self.operation.attributes: del self.operation.attributes["called_computations"] @called_computations.deleter def called_computations(self): del self.operation.attributes["called_computations"] @builtins.property def custom_call_schedule(self): if "custom_call_schedule" not in self.operation.attributes: return None return self.operation.attributes["custom_call_schedule"] @custom_call_schedule.setter def custom_call_schedule(self, value): if value is not None: self.operation.attributes["custom_call_schedule"] = value elif "custom_call_schedule" in self.operation.attributes: del self.operation.attributes["custom_call_schedule"] @custom_call_schedule.deleter def custom_call_schedule(self): del self.operation.attributes["custom_call_schedule"] @builtins.property def operand_layouts(self): if "operand_layouts" not in self.operation.attributes: return None return self.operation.attributes["operand_layouts"] @operand_layouts.setter def operand_layouts(self, value): if value is not None: self.operation.attributes["operand_layouts"] = value elif "operand_layouts" in self.operation.attributes: del self.operation.attributes["operand_layouts"] @operand_layouts.deleter def operand_layouts(self): del self.operation.attributes["operand_layouts"] @builtins.property def result_layouts(self): if "result_layouts" not in self.operation.attributes: return None return self.operation.attributes["result_layouts"] @result_layouts.setter def result_layouts(self, value): if value is not None: self.operation.attributes["result_layouts"] = value elif "result_layouts" in self.operation.attributes: del self.operation.attributes["result_layouts"] @result_layouts.deleter def result_layouts(self): del self.operation.attributes["result_layouts"] @builtins.property def output_operand_aliases(self): if "output_operand_aliases" not in self.operation.attributes: return None return self.operation.attributes["output_operand_aliases"] @output_operand_aliases.setter def output_operand_aliases(self, value): if value is not None: self.operation.attributes["output_operand_aliases"] = value elif "output_operand_aliases" in self.operation.attributes: del self.operation.attributes["output_operand_aliases"] @output_operand_aliases.deleter def output_operand_aliases(self): del self.operation.attributes["output_operand_aliases"] def custom_call(result, inputs, call_target_name, *, has_side_effect=None, backend_config=None, api_version=None, called_computations=None, custom_call_schedule=None, operand_layouts=None, result_layouts=None, output_operand_aliases=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(CustomCallOp(result=result, inputs=inputs, call_target_name=call_target_name, has_side_effect=has_side_effect, backend_config=backend_config, api_version=api_version, called_computations=called_computations, custom_call_schedule=custom_call_schedule, operand_layouts=operand_layouts, result_layouts=result_layouts, output_operand_aliases=output_operand_aliases, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class DivOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.divide" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def divide(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(DivOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class DomainOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.domain" _ODS_REGIONS = (0, True) def __init__(self, operand, kind, entry_metadata, exit_metadata, *, 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) attributes["kind"] = (kind if ( isinstance(kind, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_DomainKindAttr')) else _ods_ir.AttrBuilder.get('MHLO_DomainKindAttr')(kind, context=_ods_context)) attributes["entry_metadata"] = (entry_metadata if ( isinstance(entry_metadata, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('StrAttr')) else _ods_ir.AttrBuilder.get('StrAttr')(entry_metadata, context=_ods_context)) attributes["exit_metadata"] = (exit_metadata if ( isinstance(exit_metadata, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('StrAttr')) else _ods_ir.AttrBuilder.get('StrAttr')(exit_metadata, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @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 entry_metadata(self): return self.operation.attributes["entry_metadata"] @entry_metadata.setter def entry_metadata(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["entry_metadata"] = value @builtins.property def exit_metadata(self): return self.operation.attributes["exit_metadata"] @exit_metadata.setter def exit_metadata(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["exit_metadata"] = value @builtins.property def result(self): return self.operation.results[0] def domain(operand, kind, entry_metadata, exit_metadata, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(DomainOp(operand=operand, kind=kind, entry_metadata=entry_metadata, exit_metadata=exit_metadata, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class DotGeneralOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.dot_general" _ODS_REGIONS = (0, True) def __init__(self, result, lhs, rhs, dot_dimension_numbers, *, precision_config=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)) _ods_context = _ods_get_default_loc_context(loc) attributes["dot_dimension_numbers"] = (dot_dimension_numbers if ( isinstance(dot_dimension_numbers, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_DotDimensionNumbers')) else _ods_ir.AttrBuilder.get('MHLO_DotDimensionNumbers')(dot_dimension_numbers, context=_ods_context)) if precision_config is not None: attributes["precision_config"] = (precision_config if ( isinstance(precision_config, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_PrecisionConfigAttr')) else _ods_ir.AttrBuilder.get('MHLO_PrecisionConfigAttr')(precision_config, 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 dot_dimension_numbers(self): return self.operation.attributes["dot_dimension_numbers"] @dot_dimension_numbers.setter def dot_dimension_numbers(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["dot_dimension_numbers"] = value @builtins.property def precision_config(self): if "precision_config" not in self.operation.attributes: return None return self.operation.attributes["precision_config"] @precision_config.setter def precision_config(self, value): if value is not None: self.operation.attributes["precision_config"] = value elif "precision_config" in self.operation.attributes: del self.operation.attributes["precision_config"] @precision_config.deleter def precision_config(self): del self.operation.attributes["precision_config"] def dot_general(result, lhs, rhs, dot_dimension_numbers, *, precision_config=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(DotGeneralOp(result=result, lhs=lhs, rhs=rhs, dot_dimension_numbers=dot_dimension_numbers, precision_config=precision_config, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class DotOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.dot" _ODS_REGIONS = (0, True) def __init__(self, result, lhs, rhs, *, precision_config=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)) _ods_context = _ods_get_default_loc_context(loc) if precision_config is not None: attributes["precision_config"] = (precision_config if ( isinstance(precision_config, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_PrecisionConfigAttr')) else _ods_ir.AttrBuilder.get('MHLO_PrecisionConfigAttr')(precision_config, 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 precision_config(self): if "precision_config" not in self.operation.attributes: return None return self.operation.attributes["precision_config"] @precision_config.setter def precision_config(self, value): if value is not None: self.operation.attributes["precision_config"] = value elif "precision_config" in self.operation.attributes: del self.operation.attributes["precision_config"] @precision_config.deleter def precision_config(self): del self.operation.attributes["precision_config"] def dot(result, lhs, rhs, *, precision_config=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(DotOp(result=result, lhs=lhs, rhs=rhs, precision_config=precision_config, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class DynamicBroadcastInDimOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.dynamic_broadcast_in_dim" _ODS_REGIONS = (0, True) def __init__(self, result, operand, output_dimensions, broadcast_dimensions, *, known_expanding_dimensions=None, known_nonexpanding_dimensions=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(output_dimensions)) _ods_context = _ods_get_default_loc_context(loc) attributes["broadcast_dimensions"] = (broadcast_dimensions if ( isinstance(broadcast_dimensions, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(broadcast_dimensions, context=_ods_context)) if known_expanding_dimensions is not None: attributes["known_expanding_dimensions"] = (known_expanding_dimensions if ( isinstance(known_expanding_dimensions, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(known_expanding_dimensions, context=_ods_context)) if known_nonexpanding_dimensions is not None: attributes["known_nonexpanding_dimensions"] = (known_nonexpanding_dimensions if ( isinstance(known_nonexpanding_dimensions, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(known_nonexpanding_dimensions, 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 operand(self): return self.operation.operands[0] @builtins.property def output_dimensions(self): return self.operation.operands[1] @builtins.property def broadcast_dimensions(self): return self.operation.attributes["broadcast_dimensions"] @broadcast_dimensions.setter def broadcast_dimensions(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["broadcast_dimensions"] = value @builtins.property def known_expanding_dimensions(self): if "known_expanding_dimensions" not in self.operation.attributes: return None return self.operation.attributes["known_expanding_dimensions"] @known_expanding_dimensions.setter def known_expanding_dimensions(self, value): if value is not None: self.operation.attributes["known_expanding_dimensions"] = value elif "known_expanding_dimensions" in self.operation.attributes: del self.operation.attributes["known_expanding_dimensions"] @known_expanding_dimensions.deleter def known_expanding_dimensions(self): del self.operation.attributes["known_expanding_dimensions"] @builtins.property def known_nonexpanding_dimensions(self): if "known_nonexpanding_dimensions" not in self.operation.attributes: return None return self.operation.attributes["known_nonexpanding_dimensions"] @known_nonexpanding_dimensions.setter def known_nonexpanding_dimensions(self, value): if value is not None: self.operation.attributes["known_nonexpanding_dimensions"] = value elif "known_nonexpanding_dimensions" in self.operation.attributes: del self.operation.attributes["known_nonexpanding_dimensions"] @known_nonexpanding_dimensions.deleter def known_nonexpanding_dimensions(self): del self.operation.attributes["known_nonexpanding_dimensions"] def dynamic_broadcast_in_dim(result, operand, output_dimensions, broadcast_dimensions, *, known_expanding_dimensions=None, known_nonexpanding_dimensions=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(DynamicBroadcastInDimOp(result=result, operand=operand, output_dimensions=output_dimensions, broadcast_dimensions=broadcast_dimensions, known_expanding_dimensions=known_expanding_dimensions, known_nonexpanding_dimensions=known_nonexpanding_dimensions, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class DynamicConvOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.dynamic_conv" _ODS_REGIONS = (0, True) def __init__(self, result, lhs, rhs, d_padding, dimension_numbers, feature_group_count, batch_group_count, *, window_strides=None, padding=None, lhs_dilation=None, rhs_dilation=None, window_reversal=None, precision_config=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(d_padding)) _ods_context = _ods_get_default_loc_context(loc) if window_strides is not None: attributes["window_strides"] = (window_strides if ( isinstance(window_strides, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(window_strides, context=_ods_context)) if padding is not None: attributes["padding"] = (padding if ( isinstance(padding, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(padding, context=_ods_context)) if lhs_dilation is not None: attributes["lhs_dilation"] = (lhs_dilation if ( isinstance(lhs_dilation, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(lhs_dilation, context=_ods_context)) if rhs_dilation is not None: attributes["rhs_dilation"] = (rhs_dilation if ( isinstance(rhs_dilation, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(rhs_dilation, context=_ods_context)) if window_reversal is not None: attributes["window_reversal"] = (window_reversal if ( isinstance(window_reversal, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_BoolElementsAttr')) else _ods_ir.AttrBuilder.get('MHLO_BoolElementsAttr')(window_reversal, context=_ods_context)) attributes["dimension_numbers"] = (dimension_numbers if ( isinstance(dimension_numbers, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ConvDimensionNumbers')) else _ods_ir.AttrBuilder.get('MHLO_ConvDimensionNumbers')(dimension_numbers, context=_ods_context)) attributes["feature_group_count"] = (feature_group_count if ( isinstance(feature_group_count, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(feature_group_count, context=_ods_context)) attributes["batch_group_count"] = (batch_group_count if ( isinstance(batch_group_count, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(batch_group_count, context=_ods_context)) if precision_config is not None: attributes["precision_config"] = (precision_config if ( isinstance(precision_config, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_PrecisionConfigAttr')) else _ods_ir.AttrBuilder.get('MHLO_PrecisionConfigAttr')(precision_config, 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 d_padding(self): return self.operation.operands[2] @builtins.property def window_strides(self): if "window_strides" not in self.operation.attributes: return None return self.operation.attributes["window_strides"] @window_strides.setter def window_strides(self, value): if value is not None: self.operation.attributes["window_strides"] = value elif "window_strides" in self.operation.attributes: del self.operation.attributes["window_strides"] @window_strides.deleter def window_strides(self): del self.operation.attributes["window_strides"] @builtins.property def padding(self): if "padding" not in self.operation.attributes: return None return self.operation.attributes["padding"] @padding.setter def padding(self, value): if value is not None: self.operation.attributes["padding"] = value elif "padding" in self.operation.attributes: del self.operation.attributes["padding"] @padding.deleter def padding(self): del self.operation.attributes["padding"] @builtins.property def lhs_dilation(self): if "lhs_dilation" not in self.operation.attributes: return None return self.operation.attributes["lhs_dilation"] @lhs_dilation.setter def lhs_dilation(self, value): if value is not None: self.operation.attributes["lhs_dilation"] = value elif "lhs_dilation" in self.operation.attributes: del self.operation.attributes["lhs_dilation"] @lhs_dilation.deleter def lhs_dilation(self): del self.operation.attributes["lhs_dilation"] @builtins.property def rhs_dilation(self): if "rhs_dilation" not in self.operation.attributes: return None return self.operation.attributes["rhs_dilation"] @rhs_dilation.setter def rhs_dilation(self, value): if value is not None: self.operation.attributes["rhs_dilation"] = value elif "rhs_dilation" in self.operation.attributes: del self.operation.attributes["rhs_dilation"] @rhs_dilation.deleter def rhs_dilation(self): del self.operation.attributes["rhs_dilation"] @builtins.property def window_reversal(self): if "window_reversal" not in self.operation.attributes: return None return self.operation.attributes["window_reversal"] @window_reversal.setter def window_reversal(self, value): if value is not None: self.operation.attributes["window_reversal"] = value elif "window_reversal" in self.operation.attributes: del self.operation.attributes["window_reversal"] @window_reversal.deleter def window_reversal(self): del self.operation.attributes["window_reversal"] @builtins.property def dimension_numbers(self): return self.operation.attributes["dimension_numbers"] @dimension_numbers.setter def dimension_numbers(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["dimension_numbers"] = value @builtins.property def feature_group_count(self): return self.operation.attributes["feature_group_count"] @feature_group_count.setter def feature_group_count(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["feature_group_count"] = value @builtins.property def batch_group_count(self): return self.operation.attributes["batch_group_count"] @batch_group_count.setter def batch_group_count(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["batch_group_count"] = value @builtins.property def precision_config(self): if "precision_config" not in self.operation.attributes: return None return self.operation.attributes["precision_config"] @precision_config.setter def precision_config(self, value): if value is not None: self.operation.attributes["precision_config"] = value elif "precision_config" in self.operation.attributes: del self.operation.attributes["precision_config"] @precision_config.deleter def precision_config(self): del self.operation.attributes["precision_config"] def dynamic_conv(result, lhs, rhs, d_padding, dimension_numbers, feature_group_count, batch_group_count, *, window_strides=None, padding=None, lhs_dilation=None, rhs_dilation=None, window_reversal=None, precision_config=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(DynamicConvOp(result=result, lhs=lhs, rhs=rhs, d_padding=d_padding, dimension_numbers=dimension_numbers, feature_group_count=feature_group_count, batch_group_count=batch_group_count, window_strides=window_strides, padding=padding, lhs_dilation=lhs_dilation, rhs_dilation=rhs_dilation, window_reversal=window_reversal, precision_config=precision_config, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class DynamicGatherOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.dynamic_gather" _ODS_REGIONS = (0, True) def __init__(self, operand, start_indices, slice_sizes, dimension_numbers, *, indices_are_sorted=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(start_indices)) operands.append(_get_op_result_or_value(slice_sizes)) _ods_context = _ods_get_default_loc_context(loc) attributes["dimension_numbers"] = (dimension_numbers if ( isinstance(dimension_numbers, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_GatherDimensionNumbers')) else _ods_ir.AttrBuilder.get('MHLO_GatherDimensionNumbers')(dimension_numbers, context=_ods_context)) if indices_are_sorted is not None: attributes["indices_are_sorted"] = (indices_are_sorted if ( isinstance(indices_are_sorted, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('BoolAttr')) else _ods_ir.AttrBuilder.get('BoolAttr')(indices_are_sorted, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def start_indices(self): return self.operation.operands[1] @builtins.property def slice_sizes(self): return self.operation.operands[2] @builtins.property def dimension_numbers(self): return self.operation.attributes["dimension_numbers"] @dimension_numbers.setter def dimension_numbers(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["dimension_numbers"] = value @builtins.property def indices_are_sorted(self): if "indices_are_sorted" not in self.operation.attributes: return None return self.operation.attributes["indices_are_sorted"] @indices_are_sorted.setter def indices_are_sorted(self, value): if value is not None: self.operation.attributes["indices_are_sorted"] = value elif "indices_are_sorted" in self.operation.attributes: del self.operation.attributes["indices_are_sorted"] @indices_are_sorted.deleter def indices_are_sorted(self): del self.operation.attributes["indices_are_sorted"] def dynamic_gather(operand, start_indices, slice_sizes, dimension_numbers, *, indices_are_sorted=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(DynamicGatherOp(operand=operand, start_indices=start_indices, slice_sizes=slice_sizes, dimension_numbers=dimension_numbers, indices_are_sorted=indices_are_sorted, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class DynamicIotaOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.dynamic_iota" _ODS_REGIONS = (0, True) def __init__(self, result, output_shape, iota_dimension, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(output_shape)) _ods_context = _ods_get_default_loc_context(loc) attributes["iota_dimension"] = (iota_dimension if ( isinstance(iota_dimension, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(iota_dimension, 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 output_shape(self): return self.operation.operands[0] @builtins.property def iota_dimension(self): return self.operation.attributes["iota_dimension"] @iota_dimension.setter def iota_dimension(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["iota_dimension"] = value @builtins.property def result(self): return self.operation.results[0] def dynamic_iota(result, output_shape, iota_dimension, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(DynamicIotaOp(result=result, output_shape=output_shape, iota_dimension=iota_dimension, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class DynamicPadOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.dynamic_pad" _ODS_REGIONS = (0, True) def __init__(self, result, operand, padding_value, edge_padding_low, edge_padding_high, interior_padding, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(padding_value)) operands.append(_get_op_result_or_value(edge_padding_low)) operands.append(_get_op_result_or_value(edge_padding_high)) operands.append(_get_op_result_or_value(interior_padding)) _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 padding_value(self): return self.operation.operands[1] @builtins.property def edge_padding_low(self): return self.operation.operands[2] @builtins.property def edge_padding_high(self): return self.operation.operands[3] @builtins.property def interior_padding(self): return self.operation.operands[4] @builtins.property def result(self): return self.operation.results[0] def dynamic_pad(result, operand, padding_value, edge_padding_low, edge_padding_high, interior_padding, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(DynamicPadOp(result=result, operand=operand, padding_value=padding_value, edge_padding_low=edge_padding_low, edge_padding_high=edge_padding_high, interior_padding=interior_padding, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class DynamicReshapeOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.dynamic_reshape" _ODS_REGIONS = (0, True) def __init__(self, result, operand, output_shape, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(output_shape)) _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 output_shape(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def dynamic_reshape(result, operand, output_shape, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(DynamicReshapeOp(result=result, operand=operand, output_shape=output_shape, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class DynamicSliceOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.dynamic_slice" _ODS_REGIONS = (0, True) def __init__(self, operand, start_indices, slice_sizes, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.extend(_get_op_results_or_values(start_indices)) _ods_context = _ods_get_default_loc_context(loc) attributes["slice_sizes"] = (slice_sizes if ( isinstance(slice_sizes, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(slice_sizes, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def start_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 slice_sizes(self): return self.operation.attributes["slice_sizes"] @slice_sizes.setter def slice_sizes(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["slice_sizes"] = value @builtins.property def result(self): return self.operation.results[0] def dynamic_slice(operand, start_indices, slice_sizes, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(DynamicSliceOp(operand=operand, start_indices=start_indices, slice_sizes=slice_sizes, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class DynamicUpdateSliceOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.dynamic_update_slice" _ODS_REGIONS = (0, True) def __init__(self, operand, update, start_indices, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(update)) operands.extend(_get_op_results_or_values(start_indices)) _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 operand(self): return self.operation.operands[0] @builtins.property def update(self): return self.operation.operands[1] @builtins.property def start_indices(self): _ods_variadic_group_length = len(self.operation.operands) - 3 + 1 return self.operation.operands[2:2 + _ods_variadic_group_length] @builtins.property def result(self): return self.operation.results[0] def dynamic_update_slice(operand, update, start_indices, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(DynamicUpdateSliceOp(operand=operand, update=update, start_indices=start_indices, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class EinsumOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.einsum" _ODS_REGIONS = (0, True) def __init__(self, result, lhs, rhs, einsum_config, *, 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)) _ods_context = _ods_get_default_loc_context(loc) attributes["einsum_config"] = (einsum_config if ( isinstance(einsum_config, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('StrAttr')) else _ods_ir.AttrBuilder.get('StrAttr')(einsum_config, 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 einsum_config(self): return self.operation.attributes["einsum_config"] @einsum_config.setter def einsum_config(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["einsum_config"] = value def einsum(result, lhs, rhs, einsum_config, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(EinsumOp(result=result, lhs=lhs, rhs=rhs, einsum_config=einsum_config, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ErfOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.erf" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def erf(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ErfOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ExpOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.exponential" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def exponential(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ExpOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class Expm1Op(_ods_ir.OpView): OPERATION_NAME = "mhlo.exponential_minus_one" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def exponential_minus_one(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(Expm1Op(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class FftOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.fft" _ODS_REGIONS = (0, True) def __init__(self, operand, fft_type, fft_length, *, 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) attributes["fft_type"] = (fft_type if ( isinstance(fft_type, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_FftTypeAttr')) else _ods_ir.AttrBuilder.get('MHLO_FftTypeAttr')(fft_type, context=_ods_context)) attributes["fft_length"] = (fft_length if ( isinstance(fft_length, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(fft_length, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def fft_type(self): return self.operation.attributes["fft_type"] @fft_type.setter def fft_type(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["fft_type"] = value @builtins.property def fft_length(self): return self.operation.attributes["fft_length"] @fft_length.setter def fft_length(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["fft_length"] = value def fft(operand, fft_type, fft_length, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(FftOp(operand=operand, fft_type=fft_type, fft_length=fft_length, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class FloorOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.floor" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def floor(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(FloorOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class FusionOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.fusion" _ODS_REGIONS = (1, True) def __init__(self, results_, inputs, *, fusion_kind=None, output_operand_aliases=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(inputs)) _ods_context = _ods_get_default_loc_context(loc) if fusion_kind is not None: attributes["fusion_kind"] = (fusion_kind if ( isinstance(fusion_kind, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_FusionKindAttr')) else _ods_ir.AttrBuilder.get('MHLO_FusionKindAttr')(fusion_kind, context=_ods_context)) if output_operand_aliases is not None: attributes["output_operand_aliases"] = (output_operand_aliases if ( isinstance(output_operand_aliases, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('anonymous_850')) else _ods_ir.AttrBuilder.get('anonymous_850')(output_operand_aliases, 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 inputs(self): _ods_variadic_group_length = len(self.operation.operands) - 1 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @builtins.property def fusion_kind(self): if "fusion_kind" not in self.operation.attributes: return None return self.operation.attributes["fusion_kind"] @fusion_kind.setter def fusion_kind(self, value): if value is not None: self.operation.attributes["fusion_kind"] = value elif "fusion_kind" in self.operation.attributes: del self.operation.attributes["fusion_kind"] @fusion_kind.deleter def fusion_kind(self): del self.operation.attributes["fusion_kind"] @builtins.property def output_operand_aliases(self): if "output_operand_aliases" not in self.operation.attributes: return None return self.operation.attributes["output_operand_aliases"] @output_operand_aliases.setter def output_operand_aliases(self, value): if value is not None: self.operation.attributes["output_operand_aliases"] = value elif "output_operand_aliases" in self.operation.attributes: del self.operation.attributes["output_operand_aliases"] @output_operand_aliases.deleter def output_operand_aliases(self): del self.operation.attributes["output_operand_aliases"] @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 fused_computation(self): return self.regions[0] def fusion(results_, inputs, *, fusion_kind=None, output_operand_aliases=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(FusionOp(results_=results_, inputs=inputs, fusion_kind=fusion_kind, output_operand_aliases=output_operand_aliases, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class GatherOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.gather" _ODS_REGIONS = (0, True) def __init__(self, operand, start_indices, dimension_numbers, slice_sizes, *, indices_are_sorted=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(start_indices)) _ods_context = _ods_get_default_loc_context(loc) attributes["dimension_numbers"] = (dimension_numbers if ( isinstance(dimension_numbers, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_GatherDimensionNumbers')) else _ods_ir.AttrBuilder.get('MHLO_GatherDimensionNumbers')(dimension_numbers, context=_ods_context)) attributes["slice_sizes"] = (slice_sizes if ( isinstance(slice_sizes, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(slice_sizes, context=_ods_context)) if indices_are_sorted is not None: attributes["indices_are_sorted"] = (indices_are_sorted if ( isinstance(indices_are_sorted, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('BoolAttr')) else _ods_ir.AttrBuilder.get('BoolAttr')(indices_are_sorted, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def start_indices(self): return self.operation.operands[1] @builtins.property def dimension_numbers(self): return self.operation.attributes["dimension_numbers"] @dimension_numbers.setter def dimension_numbers(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["dimension_numbers"] = value @builtins.property def slice_sizes(self): return self.operation.attributes["slice_sizes"] @slice_sizes.setter def slice_sizes(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["slice_sizes"] = value @builtins.property def indices_are_sorted(self): if "indices_are_sorted" not in self.operation.attributes: return None return self.operation.attributes["indices_are_sorted"] @indices_are_sorted.setter def indices_are_sorted(self, value): if value is not None: self.operation.attributes["indices_are_sorted"] = value elif "indices_are_sorted" in self.operation.attributes: del self.operation.attributes["indices_are_sorted"] @indices_are_sorted.deleter def indices_are_sorted(self): del self.operation.attributes["indices_are_sorted"] def gather(operand, start_indices, dimension_numbers, slice_sizes, *, indices_are_sorted=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(GatherOp(operand=operand, start_indices=start_indices, dimension_numbers=dimension_numbers, slice_sizes=slice_sizes, indices_are_sorted=indices_are_sorted, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class GetDimensionSizeOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.get_dimension_size" _ODS_REGIONS = (0, True) def __init__(self, operand, dimension, *, 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) attributes["dimension"] = (dimension if ( isinstance(dimension, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(dimension, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(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 def get_dimension_size(operand, dimension, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(GetDimensionSizeOp(operand=operand, dimension=dimension, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class GetTupleElementOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.get_tuple_element" _ODS_REGIONS = (0, True) def __init__(self, operand, index, *, 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) attributes["index"] = (index if ( isinstance(index, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I32Attr')) else _ods_ir.AttrBuilder.get('I32Attr')(index, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(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 def get_tuple_element(operand, index, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(GetTupleElementOp(operand=operand, index=index, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class IfOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.if" _ODS_REGIONS = (2, True) def __init__(self, result, pred, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(pred)) _ods_context = _ods_get_default_loc_context(loc) results.extend(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 pred(self): return self.operation.operands[0] @builtins.property def true_branch(self): return self.regions[0] @builtins.property def false_branch(self): return self.regions[1] def if_(result, pred, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(IfOp(result=result, pred=pred, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ImagOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.imag" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def imag(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ImagOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class InfeedOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.infeed" _ODS_REGIONS = (0, True) def __init__(self, result, token, *, infeed_config=None, layout=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(token)) _ods_context = _ods_get_default_loc_context(loc) if infeed_config is not None: attributes["infeed_config"] = (infeed_config if ( isinstance(infeed_config, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('StrAttr')) else _ods_ir.AttrBuilder.get('StrAttr')(infeed_config, context=_ods_context)) if layout is not None: attributes["layout"] = (layout if ( isinstance(layout, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('ArrayAttr')) else _ods_ir.AttrBuilder.get('ArrayAttr')(layout, context=_ods_context)) results.extend(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 token(self): return self.operation.operands[0] @builtins.property def infeed_config(self): return self.operation.attributes["infeed_config"] @infeed_config.setter def infeed_config(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["infeed_config"] = value @builtins.property def layout(self): if "layout" not in self.operation.attributes: return None return self.operation.attributes["layout"] @layout.setter def layout(self, value): if value is not None: self.operation.attributes["layout"] = value elif "layout" in self.operation.attributes: del self.operation.attributes["layout"] @layout.deleter def layout(self): del self.operation.attributes["layout"] def infeed(result, token, *, infeed_config=None, layout=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(InfeedOp(result=result, token=token, infeed_config=infeed_config, layout=layout, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class IotaOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.iota" _ODS_REGIONS = (0, True) def __init__(self, output, iota_dimension, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None _ods_context = _ods_get_default_loc_context(loc) attributes["iota_dimension"] = (iota_dimension if ( isinstance(iota_dimension, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(iota_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 iota_dimension(self): return self.operation.attributes["iota_dimension"] @iota_dimension.setter def iota_dimension(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["iota_dimension"] = value @builtins.property def output(self): return self.operation.results[0] def iota(output, iota_dimension, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(IotaOp(output=output, iota_dimension=iota_dimension, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class IsFiniteOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.is_finite" _ODS_REGIONS = (0, True) def __init__(self, x, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(x)) _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 x(self): return self.operation.operands[0] @builtins.property def y(self): return self.operation.results[0] def is_finite(x, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(IsFiniteOp(x=x, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class Log1pOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.log_plus_one" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def log_plus_one(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(Log1pOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class LogOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.log" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def log(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(LogOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class LogisticOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.logistic" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def logistic(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(LogisticOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class MapOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.map" _ODS_REGIONS = (1, True) def __init__(self, result, inputs, dimensions, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(inputs)) _ods_context = _ods_get_default_loc_context(loc) attributes["dimensions"] = (dimensions if ( isinstance(dimensions, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(dimensions, 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 inputs(self): _ods_variadic_group_length = len(self.operation.operands) - 1 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @builtins.property def dimensions(self): return self.operation.attributes["dimensions"] @dimensions.setter def dimensions(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["dimensions"] = value @builtins.property def computation(self): return self.regions[0] def map(result, inputs, dimensions, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(MapOp(result=result, inputs=inputs, dimensions=dimensions, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class MaxOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.maximum" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def maximum(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(MaxOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class MinOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.minimum" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def minimum(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(MinOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class MinimumBroadcastShapesOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.minimum_broadcast_shapes" _ODS_REGIONS = (0, True) def __init__(self, results_, shapes, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(shapes)) _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 shapes(self): _ods_variadic_group_length = len(self.operation.operands) - 1 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @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] def minimum_broadcast_shapes(results_, shapes, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(MinimumBroadcastShapesOp(results_=results_, shapes=shapes, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class MulOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.multiply" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def multiply(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(MulOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class NegOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.negate" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def negate(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(NegOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class NotOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.not" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def not_(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(NotOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class OptimizationBarrierOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.optimization_barrier" _ODS_REGIONS = (0, True) def __init__(self, operand, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(operand)) _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 operand(self): _ods_variadic_group_length = len(self.operation.operands) - 1 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @builtins.property def result(self): _ods_variadic_group_length = len(self.operation.results) - 1 + 1 return self.operation.results[0:0 + _ods_variadic_group_length] def optimization_barrier(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(OptimizationBarrierOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class OrOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.or" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def or_(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(OrOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class OutfeedOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.outfeed" _ODS_REGIONS = (0, True) def __init__(self, inputs, token, *, outfeed_config=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(inputs)) operands.append(_get_op_result_or_value(token)) _ods_context = _ods_get_default_loc_context(loc) if outfeed_config is not None: attributes["outfeed_config"] = (outfeed_config if ( isinstance(outfeed_config, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('StrAttr')) else _ods_ir.AttrBuilder.get('StrAttr')(outfeed_config, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def inputs(self): _ods_variadic_group_length = len(self.operation.operands) - 2 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @builtins.property def token(self): _ods_variadic_group_length = len(self.operation.operands) - 2 + 1 return self.operation.operands[1 + _ods_variadic_group_length - 1] @builtins.property def outfeed_config(self): return self.operation.attributes["outfeed_config"] @outfeed_config.setter def outfeed_config(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["outfeed_config"] = value def outfeed(inputs, token, *, outfeed_config=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(OutfeedOp(inputs=inputs, token=token, outfeed_config=outfeed_config, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class PadOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.pad" _ODS_REGIONS = (0, True) def __init__(self, operand, padding_value, edge_padding_low, edge_padding_high, interior_padding, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(padding_value)) _ods_context = _ods_get_default_loc_context(loc) attributes["edge_padding_low"] = (edge_padding_low if ( isinstance(edge_padding_low, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(edge_padding_low, context=_ods_context)) attributes["edge_padding_high"] = (edge_padding_high if ( isinstance(edge_padding_high, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(edge_padding_high, context=_ods_context)) attributes["interior_padding"] = (interior_padding if ( isinstance(interior_padding, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(interior_padding, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def padding_value(self): return self.operation.operands[1] @builtins.property def edge_padding_low(self): return self.operation.attributes["edge_padding_low"] @edge_padding_low.setter def edge_padding_low(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["edge_padding_low"] = value @builtins.property def edge_padding_high(self): return self.operation.attributes["edge_padding_high"] @edge_padding_high.setter def edge_padding_high(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["edge_padding_high"] = value @builtins.property def interior_padding(self): return self.operation.attributes["interior_padding"] @interior_padding.setter def interior_padding(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["interior_padding"] = value def pad(operand, padding_value, edge_padding_low, edge_padding_high, interior_padding, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(PadOp(operand=operand, padding_value=padding_value, edge_padding_low=edge_padding_low, edge_padding_high=edge_padding_high, interior_padding=interior_padding, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class PartitionIdOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.partition_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)) def partition_id(*, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(PartitionIdOp(loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class PopulationCountOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.popcnt" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def popcnt(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(PopulationCountOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class PowOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.power" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def power(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(PowOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class RealDynamicSliceOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.real_dynamic_slice" _ODS_REGIONS = (0, True) def __init__(self, result, operand, start_indices, limit_indices, strides, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(start_indices)) operands.append(_get_op_result_or_value(limit_indices)) operands.append(_get_op_result_or_value(strides)) _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 start_indices(self): return self.operation.operands[1] @builtins.property def limit_indices(self): return self.operation.operands[2] @builtins.property def strides(self): return self.operation.operands[3] @builtins.property def result(self): return self.operation.results[0] def real_dynamic_slice(result, operand, start_indices, limit_indices, strides, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(RealDynamicSliceOp(result=result, operand=operand, start_indices=start_indices, limit_indices=limit_indices, strides=strides, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class RealOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.real" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def real(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(RealOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class RecvOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.recv" _ODS_REGIONS = (0, True) def __init__(self, result, token, channel_handle, *, is_host_transfer=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(token)) _ods_context = _ods_get_default_loc_context(loc) attributes["channel_handle"] = (channel_handle if ( isinstance(channel_handle, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ChannelHandle')) else _ods_ir.AttrBuilder.get('MHLO_ChannelHandle')(channel_handle, context=_ods_context)) if is_host_transfer is not None: attributes["is_host_transfer"] = (is_host_transfer if ( isinstance(is_host_transfer, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('BoolAttr')) else _ods_ir.AttrBuilder.get('BoolAttr')(is_host_transfer, context=_ods_context)) results.extend(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 token(self): return self.operation.operands[0] @builtins.property def channel_handle(self): return self.operation.attributes["channel_handle"] @channel_handle.setter def channel_handle(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["channel_handle"] = value @builtins.property def is_host_transfer(self): if "is_host_transfer" not in self.operation.attributes: return None return self.operation.attributes["is_host_transfer"] @is_host_transfer.setter def is_host_transfer(self, value): if value is not None: self.operation.attributes["is_host_transfer"] = value elif "is_host_transfer" in self.operation.attributes: del self.operation.attributes["is_host_transfer"] @is_host_transfer.deleter def is_host_transfer(self): del self.operation.attributes["is_host_transfer"] def recv(result, token, channel_handle, *, is_host_transfer=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(RecvOp(result=result, token=token, channel_handle=channel_handle, is_host_transfer=is_host_transfer, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ReduceOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.reduce" _ODS_REGIONS = (1, True) def __init__(self, result, inputs, init_values, dimensions, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(inputs)) operands.extend(_get_op_results_or_values(init_values)) _ods_context = _ods_get_default_loc_context(loc) attributes["dimensions"] = (dimensions if ( isinstance(dimensions, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(dimensions, context=_ods_context)) results.extend(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 inputs(self): start, pg = _ods_equally_sized_accessor(operation.operands, 2, 0, 0) return self.operation.operands[start:start + pg] @builtins.property def init_values(self): start, pg = _ods_equally_sized_accessor(operation.operands, 2, 0, 1) return self.operation.operands[start:start + pg] @builtins.property def dimensions(self): return self.operation.attributes["dimensions"] @dimensions.setter def dimensions(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["dimensions"] = value @builtins.property def body(self): return self.regions[0] def reduce(result, inputs, init_values, dimensions, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ReduceOp(result=result, inputs=inputs, init_values=init_values, dimensions=dimensions, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ReducePrecisionOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.reduce_precision" _ODS_REGIONS = (0, True) def __init__(self, operand, exponent_bits, mantissa_bits, *, 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) attributes["exponent_bits"] = (exponent_bits if ( isinstance(exponent_bits, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I32Attr')) else _ods_ir.AttrBuilder.get('I32Attr')(exponent_bits, context=_ods_context)) attributes["mantissa_bits"] = (mantissa_bits if ( isinstance(mantissa_bits, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I32Attr')) else _ods_ir.AttrBuilder.get('I32Attr')(mantissa_bits, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def exponent_bits(self): return self.operation.attributes["exponent_bits"] @exponent_bits.setter def exponent_bits(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["exponent_bits"] = value @builtins.property def mantissa_bits(self): return self.operation.attributes["mantissa_bits"] @mantissa_bits.setter def mantissa_bits(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["mantissa_bits"] = value @builtins.property def output(self): return self.operation.results[0] def reduce_precision(operand, exponent_bits, mantissa_bits, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ReducePrecisionOp(operand=operand, exponent_bits=exponent_bits, mantissa_bits=mantissa_bits, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ReduceScatterOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.reduce_scatter" _ODS_REGIONS = (1, True) def __init__(self, result, operand, scatter_dimension, replica_groups, *, channel_handle=None, use_global_device_ids=None, 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) attributes["scatter_dimension"] = (scatter_dimension if ( isinstance(scatter_dimension, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(scatter_dimension, context=_ods_context)) attributes["replica_groups"] = (replica_groups if ( isinstance(replica_groups, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(replica_groups, context=_ods_context)) if channel_handle is not None: attributes["channel_handle"] = (channel_handle if ( isinstance(channel_handle, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ChannelHandle')) else _ods_ir.AttrBuilder.get('MHLO_ChannelHandle')(channel_handle, context=_ods_context)) if bool(use_global_device_ids): attributes["use_global_device_ids"] = _ods_ir.UnitAttr.get( _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 scatter_dimension(self): return self.operation.attributes["scatter_dimension"] @scatter_dimension.setter def scatter_dimension(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["scatter_dimension"] = value @builtins.property def replica_groups(self): return self.operation.attributes["replica_groups"] @replica_groups.setter def replica_groups(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["replica_groups"] = value @builtins.property def channel_handle(self): if "channel_handle" not in self.operation.attributes: return None return self.operation.attributes["channel_handle"] @channel_handle.setter def channel_handle(self, value): if value is not None: self.operation.attributes["channel_handle"] = value elif "channel_handle" in self.operation.attributes: del self.operation.attributes["channel_handle"] @channel_handle.deleter def channel_handle(self): del self.operation.attributes["channel_handle"] @builtins.property def use_global_device_ids(self): return "use_global_device_ids" in self.operation.attributes @use_global_device_ids.setter def use_global_device_ids(self, value): if bool(value): self.operation.attributes["use_global_device_ids"] = _ods_ir.UnitAttr.get() elif "use_global_device_ids" in self.operation.attributes: del self.operation.attributes["use_global_device_ids"] @use_global_device_ids.deleter def use_global_device_ids(self): del self.operation.attributes["use_global_device_ids"] @builtins.property def computation(self): return self.regions[0] def reduce_scatter(result, operand, scatter_dimension, replica_groups, *, channel_handle=None, use_global_device_ids=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ReduceScatterOp(result=result, operand=operand, scatter_dimension=scatter_dimension, replica_groups=replica_groups, channel_handle=channel_handle, use_global_device_ids=use_global_device_ids, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ReduceWindowOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.reduce_window" _ODS_REGIONS = (1, True) def __init__(self, result, inputs, init_values, window_dimensions, *, window_strides=None, base_dilations=None, window_dilations=None, padding=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(inputs)) operands.extend(_get_op_results_or_values(init_values)) _ods_context = _ods_get_default_loc_context(loc) attributes["window_dimensions"] = (window_dimensions if ( isinstance(window_dimensions, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(window_dimensions, context=_ods_context)) if window_strides is not None: attributes["window_strides"] = (window_strides if ( isinstance(window_strides, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(window_strides, context=_ods_context)) if base_dilations is not None: attributes["base_dilations"] = (base_dilations if ( isinstance(base_dilations, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(base_dilations, context=_ods_context)) if window_dilations is not None: attributes["window_dilations"] = (window_dilations if ( isinstance(window_dilations, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(window_dilations, context=_ods_context)) if padding is not None: attributes["padding"] = (padding if ( isinstance(padding, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(padding, context=_ods_context)) results.extend(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 inputs(self): start, pg = _ods_equally_sized_accessor(operation.operands, 2, 0, 0) return self.operation.operands[start:start + pg] @builtins.property def init_values(self): start, pg = _ods_equally_sized_accessor(operation.operands, 2, 0, 1) return self.operation.operands[start:start + pg] @builtins.property def window_dimensions(self): return self.operation.attributes["window_dimensions"] @window_dimensions.setter def window_dimensions(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["window_dimensions"] = value @builtins.property def window_strides(self): if "window_strides" not in self.operation.attributes: return None return self.operation.attributes["window_strides"] @window_strides.setter def window_strides(self, value): if value is not None: self.operation.attributes["window_strides"] = value elif "window_strides" in self.operation.attributes: del self.operation.attributes["window_strides"] @window_strides.deleter def window_strides(self): del self.operation.attributes["window_strides"] @builtins.property def base_dilations(self): if "base_dilations" not in self.operation.attributes: return None return self.operation.attributes["base_dilations"] @base_dilations.setter def base_dilations(self, value): if value is not None: self.operation.attributes["base_dilations"] = value elif "base_dilations" in self.operation.attributes: del self.operation.attributes["base_dilations"] @base_dilations.deleter def base_dilations(self): del self.operation.attributes["base_dilations"] @builtins.property def window_dilations(self): if "window_dilations" not in self.operation.attributes: return None return self.operation.attributes["window_dilations"] @window_dilations.setter def window_dilations(self, value): if value is not None: self.operation.attributes["window_dilations"] = value elif "window_dilations" in self.operation.attributes: del self.operation.attributes["window_dilations"] @window_dilations.deleter def window_dilations(self): del self.operation.attributes["window_dilations"] @builtins.property def padding(self): if "padding" not in self.operation.attributes: return None return self.operation.attributes["padding"] @padding.setter def padding(self, value): if value is not None: self.operation.attributes["padding"] = value elif "padding" in self.operation.attributes: del self.operation.attributes["padding"] @padding.deleter def padding(self): del self.operation.attributes["padding"] @builtins.property def body(self): return self.regions[0] def reduce_window(result, inputs, init_values, window_dimensions, *, window_strides=None, base_dilations=None, window_dilations=None, padding=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ReduceWindowOp(result=result, inputs=inputs, init_values=init_values, window_dimensions=window_dimensions, window_strides=window_strides, base_dilations=base_dilations, window_dilations=window_dilations, padding=padding, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class RemOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.remainder" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def remainder(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(RemOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ReplicaIdOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.replica_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)) def replica_id(*, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ReplicaIdOp(loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ReshapeOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.reshape" _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] def reshape(result, operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ReshapeOp(result=result, operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ReturnOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.return" _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] def return_(results_, *, loc=None, ip=None) -> _ods_ir.Operation: return _get_op_result_or_op_results(ReturnOp(results_=results_, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ReverseOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.reverse" _ODS_REGIONS = (0, True) def __init__(self, operand, dimensions, *, 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) attributes["dimensions"] = (dimensions if ( isinstance(dimensions, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(dimensions, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def dimensions(self): return self.operation.attributes["dimensions"] @dimensions.setter def dimensions(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["dimensions"] = value def reverse(operand, dimensions, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ReverseOp(operand=operand, dimensions=dimensions, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class RngBitGeneratorOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.rng_bit_generator" _ODS_REGIONS = (0, True) def __init__(self, output_state, output, rng_algorithm, initial_state, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(initial_state)) _ods_context = _ods_get_default_loc_context(loc) attributes["rng_algorithm"] = (rng_algorithm if ( isinstance(rng_algorithm, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_RngAlgorithmAttr')) else _ods_ir.AttrBuilder.get('MHLO_RngAlgorithmAttr')(rng_algorithm, context=_ods_context)) results.append(output_state) 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 initial_state(self): return self.operation.operands[0] @builtins.property def rng_algorithm(self): return self.operation.attributes["rng_algorithm"] @rng_algorithm.setter def rng_algorithm(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["rng_algorithm"] = value @builtins.property def output_state(self): return self.operation.results[0] @builtins.property def output(self): return self.operation.results[1] def rng_bit_generator(output_state, output, rng_algorithm, initial_state, *, loc=None, ip=None) -> _Sequence[_ods_ir.Value]: return _get_op_result_or_op_results(RngBitGeneratorOp(output_state=output_state, output=output, rng_algorithm=rng_algorithm, initial_state=initial_state, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class RngOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.rng" _ODS_REGIONS = (0, True) def __init__(self, a, b, shape, rng_distribution, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(a)) operands.append(_get_op_result_or_value(b)) operands.append(_get_op_result_or_value(shape)) _ods_context = _ods_get_default_loc_context(loc) attributes["rng_distribution"] = (rng_distribution if ( isinstance(rng_distribution, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_RngDistributionAttr')) else _ods_ir.AttrBuilder.get('MHLO_RngDistributionAttr')(rng_distribution, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def a(self): return self.operation.operands[0] @builtins.property def b(self): return self.operation.operands[1] @builtins.property def shape(self): return self.operation.operands[2] @builtins.property def rng_distribution(self): return self.operation.attributes["rng_distribution"] @rng_distribution.setter def rng_distribution(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["rng_distribution"] = value @builtins.property def result(self): return self.operation.results[0] def rng(a, b, shape, rng_distribution, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(RngOp(a=a, b=b, shape=shape, rng_distribution=rng_distribution, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class RoundNearestEvenOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.round_nearest_even" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def round_nearest_even(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(RoundNearestEvenOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class RoundOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.round_nearest_afz" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def round_nearest_afz(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(RoundOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class RsqrtOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.rsqrt" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def rsqrt(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(RsqrtOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ScatterOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.scatter" _ODS_REGIONS = (1, True) def __init__(self, result, inputs, scatter_indices, updates, scatter_dimension_numbers, *, indices_are_sorted=None, unique_indices=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(inputs)) operands.append(_get_op_result_or_value(scatter_indices)) operands.extend(_get_op_results_or_values(updates)) _ods_context = _ods_get_default_loc_context(loc) attributes["scatter_dimension_numbers"] = (scatter_dimension_numbers if ( isinstance(scatter_dimension_numbers, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ScatterDimensionNumbers')) else _ods_ir.AttrBuilder.get('MHLO_ScatterDimensionNumbers')(scatter_dimension_numbers, context=_ods_context)) if indices_are_sorted is not None: attributes["indices_are_sorted"] = (indices_are_sorted if ( isinstance(indices_are_sorted, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('BoolAttr')) else _ods_ir.AttrBuilder.get('BoolAttr')(indices_are_sorted, context=_ods_context)) if unique_indices is not None: attributes["unique_indices"] = (unique_indices if ( isinstance(unique_indices, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('BoolAttr')) else _ods_ir.AttrBuilder.get('BoolAttr')(unique_indices, context=_ods_context)) results.extend(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 inputs(self): start, pg = _ods_equally_sized_accessor(operation.operands, 2, 0, 0) return self.operation.operands[start:start + pg] @builtins.property def scatter_indices(self): start, pg = _ods_equally_sized_accessor(operation.operands, 2, 0, 1) return self.operation.operands[start] @builtins.property def updates(self): start, pg = _ods_equally_sized_accessor(operation.operands, 2, 1, 1) return self.operation.operands[start:start + pg] @builtins.property def scatter_dimension_numbers(self): return self.operation.attributes["scatter_dimension_numbers"] @scatter_dimension_numbers.setter def scatter_dimension_numbers(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["scatter_dimension_numbers"] = value @builtins.property def indices_are_sorted(self): if "indices_are_sorted" not in self.operation.attributes: return None return self.operation.attributes["indices_are_sorted"] @indices_are_sorted.setter def indices_are_sorted(self, value): if value is not None: self.operation.attributes["indices_are_sorted"] = value elif "indices_are_sorted" in self.operation.attributes: del self.operation.attributes["indices_are_sorted"] @indices_are_sorted.deleter def indices_are_sorted(self): del self.operation.attributes["indices_are_sorted"] @builtins.property def unique_indices(self): if "unique_indices" not in self.operation.attributes: return None return self.operation.attributes["unique_indices"] @unique_indices.setter def unique_indices(self, value): if value is not None: self.operation.attributes["unique_indices"] = value elif "unique_indices" in self.operation.attributes: del self.operation.attributes["unique_indices"] @unique_indices.deleter def unique_indices(self): del self.operation.attributes["unique_indices"] @builtins.property def update_computation(self): return self.regions[0] def scatter(result, inputs, scatter_indices, updates, scatter_dimension_numbers, *, indices_are_sorted=None, unique_indices=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ScatterOp(result=result, inputs=inputs, scatter_indices=scatter_indices, updates=updates, scatter_dimension_numbers=scatter_dimension_numbers, indices_are_sorted=indices_are_sorted, unique_indices=unique_indices, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class SelectAndScatterOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.select_and_scatter" _ODS_REGIONS = (2, True) def __init__(self, result, operand, source, init_value, *, window_dimensions=None, window_strides=None, padding=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(source)) operands.append(_get_op_result_or_value(init_value)) _ods_context = _ods_get_default_loc_context(loc) if window_dimensions is not None: attributes["window_dimensions"] = (window_dimensions if ( isinstance(window_dimensions, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(window_dimensions, context=_ods_context)) if window_strides is not None: attributes["window_strides"] = (window_strides if ( isinstance(window_strides, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(window_strides, context=_ods_context)) if padding is not None: attributes["padding"] = (padding if ( isinstance(padding, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(padding, 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 operand(self): return self.operation.operands[0] @builtins.property def source(self): return self.operation.operands[1] @builtins.property def init_value(self): return self.operation.operands[2] @builtins.property def window_dimensions(self): if "window_dimensions" not in self.operation.attributes: return None return self.operation.attributes["window_dimensions"] @window_dimensions.setter def window_dimensions(self, value): if value is not None: self.operation.attributes["window_dimensions"] = value elif "window_dimensions" in self.operation.attributes: del self.operation.attributes["window_dimensions"] @window_dimensions.deleter def window_dimensions(self): del self.operation.attributes["window_dimensions"] @builtins.property def window_strides(self): if "window_strides" not in self.operation.attributes: return None return self.operation.attributes["window_strides"] @window_strides.setter def window_strides(self, value): if value is not None: self.operation.attributes["window_strides"] = value elif "window_strides" in self.operation.attributes: del self.operation.attributes["window_strides"] @window_strides.deleter def window_strides(self): del self.operation.attributes["window_strides"] @builtins.property def padding(self): if "padding" not in self.operation.attributes: return None return self.operation.attributes["padding"] @padding.setter def padding(self, value): if value is not None: self.operation.attributes["padding"] = value elif "padding" in self.operation.attributes: del self.operation.attributes["padding"] @padding.deleter def padding(self): del self.operation.attributes["padding"] @builtins.property def select(self): return self.regions[0] @builtins.property def scatter(self): return self.regions[1] def select_and_scatter(result, operand, source, init_value, *, window_dimensions=None, window_strides=None, padding=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(SelectAndScatterOp(result=result, operand=operand, source=source, init_value=init_value, window_dimensions=window_dimensions, window_strides=window_strides, padding=padding, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class SelectOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.select" _ODS_REGIONS = (0, True) def __init__(self, pred, on_true, on_false, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(pred)) operands.append(_get_op_result_or_value(on_true)) operands.append(_get_op_result_or_value(on_false)) _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 pred(self): return self.operation.operands[0] @builtins.property def on_true(self): return self.operation.operands[1] @builtins.property def on_false(self): return self.operation.operands[2] @builtins.property def result(self): return self.operation.results[0] def select(pred, on_true, on_false, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(SelectOp(pred=pred, on_true=on_true, on_false=on_false, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class SendOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.send" _ODS_REGIONS = (0, True) def __init__(self, inputs, token, channel_handle, *, is_host_transfer=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(inputs)) operands.append(_get_op_result_or_value(token)) _ods_context = _ods_get_default_loc_context(loc) attributes["channel_handle"] = (channel_handle if ( isinstance(channel_handle, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_ChannelHandle')) else _ods_ir.AttrBuilder.get('MHLO_ChannelHandle')(channel_handle, context=_ods_context)) if is_host_transfer is not None: attributes["is_host_transfer"] = (is_host_transfer if ( isinstance(is_host_transfer, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('BoolAttr')) else _ods_ir.AttrBuilder.get('BoolAttr')(is_host_transfer, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def inputs(self): _ods_variadic_group_length = len(self.operation.operands) - 2 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @builtins.property def token(self): _ods_variadic_group_length = len(self.operation.operands) - 2 + 1 return self.operation.operands[1 + _ods_variadic_group_length - 1] @builtins.property def channel_handle(self): return self.operation.attributes["channel_handle"] @channel_handle.setter def channel_handle(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["channel_handle"] = value @builtins.property def is_host_transfer(self): if "is_host_transfer" not in self.operation.attributes: return None return self.operation.attributes["is_host_transfer"] @is_host_transfer.setter def is_host_transfer(self, value): if value is not None: self.operation.attributes["is_host_transfer"] = value elif "is_host_transfer" in self.operation.attributes: del self.operation.attributes["is_host_transfer"] @is_host_transfer.deleter def is_host_transfer(self): del self.operation.attributes["is_host_transfer"] def send(inputs, token, channel_handle, *, is_host_transfer=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(SendOp(inputs=inputs, token=token, channel_handle=channel_handle, is_host_transfer=is_host_transfer, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class SetDimensionSizeOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.set_dimension_size" _ODS_REGIONS = (0, True) def __init__(self, operand, size, dimension, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(size)) _ods_context = _ods_get_default_loc_context(loc) attributes["dimension"] = (dimension if ( isinstance(dimension, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(dimension, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def size(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 def set_dimension_size(operand, size, dimension, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(SetDimensionSizeOp(operand=operand, size=size, dimension=dimension, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ShiftLeftOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.shift_left" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def shift_left(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ShiftLeftOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ShiftRightArithmeticOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.shift_right_arithmetic" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def shift_right_arithmetic(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ShiftRightArithmeticOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class ShiftRightLogicalOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.shift_right_logical" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def shift_right_logical(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(ShiftRightLogicalOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class SignOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.sign" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def sign(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(SignOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class SineOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.sine" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def sine(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(SineOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class SliceOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.slice" _ODS_REGIONS = (0, True) def __init__(self, operand, start_indices, limit_indices, strides, *, 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) attributes["start_indices"] = (start_indices if ( isinstance(start_indices, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(start_indices, context=_ods_context)) attributes["limit_indices"] = (limit_indices if ( isinstance(limit_indices, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(limit_indices, context=_ods_context)) attributes["strides"] = (strides if ( isinstance(strides, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(strides, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def start_indices(self): return self.operation.attributes["start_indices"] @start_indices.setter def start_indices(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["start_indices"] = value @builtins.property def limit_indices(self): return self.operation.attributes["limit_indices"] @limit_indices.setter def limit_indices(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["limit_indices"] = value @builtins.property def strides(self): return self.operation.attributes["strides"] @strides.setter def strides(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["strides"] = value def slice(operand, start_indices, limit_indices, strides, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(SliceOp(operand=operand, start_indices=start_indices, limit_indices=limit_indices, strides=strides, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class SortOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.sort" _ODS_REGIONS = (1, True) def __init__(self, result, inputs, *, dimension=None, is_stable=None, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(inputs)) _ods_context = _ods_get_default_loc_context(loc) if dimension is not None: attributes["dimension"] = (dimension if ( isinstance(dimension, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(dimension, context=_ods_context)) if is_stable is not None: attributes["is_stable"] = (is_stable if ( isinstance(is_stable, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('BoolAttr')) else _ods_ir.AttrBuilder.get('BoolAttr')(is_stable, context=_ods_context)) results.extend(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 inputs(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): 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 is_stable(self): if "is_stable" not in self.operation.attributes: return None return self.operation.attributes["is_stable"] @is_stable.setter def is_stable(self, value): if value is not None: self.operation.attributes["is_stable"] = value elif "is_stable" in self.operation.attributes: del self.operation.attributes["is_stable"] @is_stable.deleter def is_stable(self): del self.operation.attributes["is_stable"] @builtins.property def comparator(self): return self.regions[0] def sort(result, inputs, *, dimension=None, is_stable=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(SortOp(result=result, inputs=inputs, dimension=dimension, is_stable=is_stable, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class SparseDotOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.sparse_dot" _ODS_REGIONS = (0, True) def __init__(self, result, lhs, rhs, meta, dot_dimension_numbers, *, lhs_sparsity=None, rhs_sparsity=None, precision_config=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.extend(_get_op_results_or_values(meta)) _ods_context = _ods_get_default_loc_context(loc) if lhs_sparsity is not None: attributes["lhs_sparsity"] = (lhs_sparsity if ( isinstance(lhs_sparsity, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_SparsityDescriptor')) else _ods_ir.AttrBuilder.get('MHLO_SparsityDescriptor')(lhs_sparsity, context=_ods_context)) if rhs_sparsity is not None: attributes["rhs_sparsity"] = (rhs_sparsity if ( isinstance(rhs_sparsity, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_SparsityDescriptor')) else _ods_ir.AttrBuilder.get('MHLO_SparsityDescriptor')(rhs_sparsity, context=_ods_context)) attributes["dot_dimension_numbers"] = (dot_dimension_numbers if ( isinstance(dot_dimension_numbers, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_DotDimensionNumbers')) else _ods_ir.AttrBuilder.get('MHLO_DotDimensionNumbers')(dot_dimension_numbers, context=_ods_context)) if precision_config is not None: attributes["precision_config"] = (precision_config if ( isinstance(precision_config, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_PrecisionConfigAttr')) else _ods_ir.AttrBuilder.get('MHLO_PrecisionConfigAttr')(precision_config, 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 meta(self): _ods_variadic_group_length = len(self.operation.operands) - 3 + 1 return self.operation.operands[2:2 + _ods_variadic_group_length] @builtins.property def lhs_sparsity(self): if "lhs_sparsity" not in self.operation.attributes: return None return self.operation.attributes["lhs_sparsity"] @lhs_sparsity.setter def lhs_sparsity(self, value): if value is not None: self.operation.attributes["lhs_sparsity"] = value elif "lhs_sparsity" in self.operation.attributes: del self.operation.attributes["lhs_sparsity"] @lhs_sparsity.deleter def lhs_sparsity(self): del self.operation.attributes["lhs_sparsity"] @builtins.property def rhs_sparsity(self): if "rhs_sparsity" not in self.operation.attributes: return None return self.operation.attributes["rhs_sparsity"] @rhs_sparsity.setter def rhs_sparsity(self, value): if value is not None: self.operation.attributes["rhs_sparsity"] = value elif "rhs_sparsity" in self.operation.attributes: del self.operation.attributes["rhs_sparsity"] @rhs_sparsity.deleter def rhs_sparsity(self): del self.operation.attributes["rhs_sparsity"] @builtins.property def dot_dimension_numbers(self): return self.operation.attributes["dot_dimension_numbers"] @dot_dimension_numbers.setter def dot_dimension_numbers(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["dot_dimension_numbers"] = value @builtins.property def precision_config(self): if "precision_config" not in self.operation.attributes: return None return self.operation.attributes["precision_config"] @precision_config.setter def precision_config(self, value): if value is not None: self.operation.attributes["precision_config"] = value elif "precision_config" in self.operation.attributes: del self.operation.attributes["precision_config"] @precision_config.deleter def precision_config(self): del self.operation.attributes["precision_config"] def sparse_dot(result, lhs, rhs, meta, dot_dimension_numbers, *, lhs_sparsity=None, rhs_sparsity=None, precision_config=None, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(SparseDotOp(result=result, lhs=lhs, rhs=rhs, meta=meta, dot_dimension_numbers=dot_dimension_numbers, lhs_sparsity=lhs_sparsity, rhs_sparsity=rhs_sparsity, precision_config=precision_config, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class SqrtOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.sqrt" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def sqrt(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(SqrtOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class StochasticConvertOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.stochastic_convert" _ODS_REGIONS = (0, True) def __init__(self, result, operand, random, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(random)) _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 random(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def stochastic_convert(result, operand, random, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(StochasticConvertOp(result=result, operand=operand, random=random, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class SubtractOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.subtract" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def subtract(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(SubtractOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class TanOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.tan" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def tan(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(TanOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class TanhOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.tanh" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def tanh(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(TanhOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class TopKOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.topk" _ODS_REGIONS = (0, True) def __init__(self, operand, k, *, largest=None, 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) attributes["k"] = (k if ( isinstance(k, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(k, context=_ods_context)) if largest is not None: attributes["largest"] = (largest if ( isinstance(largest, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('BoolAttr')) else _ods_ir.AttrBuilder.get('BoolAttr')(largest, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def k(self): return self.operation.attributes["k"] @k.setter def k(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["k"] = value @builtins.property def largest(self): if "largest" not in self.operation.attributes: return None return self.operation.attributes["largest"] @largest.setter def largest(self, value): if value is not None: self.operation.attributes["largest"] = value elif "largest" in self.operation.attributes: del self.operation.attributes["largest"] @largest.deleter def largest(self): del self.operation.attributes["largest"] @builtins.property def values(self): return self.operation.results[0] @builtins.property def indices(self): return self.operation.results[1] def topk(operand, k, *, largest=None, loc=None, ip=None) -> _Sequence[_ods_ir.Value]: return _get_op_result_or_op_results(TopKOp(operand=operand, k=k, largest=largest, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class TorchIndexSelectOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.torch_index_select" _ODS_REGIONS = (0, True) def __init__(self, result, operand, index, dim, batch_dims, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(operand)) operands.append(_get_op_result_or_value(index)) _ods_context = _ods_get_default_loc_context(loc) attributes["dim"] = (dim if ( isinstance(dim, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(dim, context=_ods_context)) attributes["batch_dims"] = (batch_dims if ( isinstance(batch_dims, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(batch_dims, 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 operand(self): return self.operation.operands[0] @builtins.property def index(self): return self.operation.operands[1] @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 batch_dims(self): return self.operation.attributes["batch_dims"] @batch_dims.setter def batch_dims(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["batch_dims"] = value def torch_index_select(result, operand, index, dim, batch_dims, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(TorchIndexSelectOp(result=result, operand=operand, index=index, dim=dim, batch_dims=batch_dims, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class TraceOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.trace" _ODS_REGIONS = (0, True) def __init__(self, operand, tag, *, 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) attributes["tag"] = (tag if ( isinstance(tag, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('StrAttr')) else _ods_ir.AttrBuilder.get('StrAttr')(tag, 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 operand(self): return self.operation.operands[0] @builtins.property def tag(self): return self.operation.attributes["tag"] @tag.setter def tag(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["tag"] = value def trace(operand, tag, *, loc=None, ip=None) -> _ods_ir.Operation: return _get_op_result_or_op_results(TraceOp(operand=operand, tag=tag, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class TransposeOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.transpose" _ODS_REGIONS = (0, True) def __init__(self, operand, permutation, *, 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) attributes["permutation"] = (permutation if ( isinstance(permutation, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64ElementsAttr')) else _ods_ir.AttrBuilder.get('I64ElementsAttr')(permutation, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def operand(self): return self.operation.operands[0] @builtins.property def permutation(self): return self.operation.attributes["permutation"] @permutation.setter def permutation(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["permutation"] = value def transpose(operand, permutation, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(TransposeOp(operand=operand, permutation=permutation, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class TriangularSolveOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.triangular_solve" _ODS_REGIONS = (0, True) def __init__(self, a, b, left_side, lower, unit_diagonal, transpose_a, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.append(_get_op_result_or_value(a)) operands.append(_get_op_result_or_value(b)) _ods_context = _ods_get_default_loc_context(loc) attributes["left_side"] = (left_side if ( isinstance(left_side, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('BoolAttr')) else _ods_ir.AttrBuilder.get('BoolAttr')(left_side, context=_ods_context)) attributes["lower"] = (lower if ( isinstance(lower, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('BoolAttr')) else _ods_ir.AttrBuilder.get('BoolAttr')(lower, context=_ods_context)) attributes["unit_diagonal"] = (unit_diagonal if ( isinstance(unit_diagonal, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('BoolAttr')) else _ods_ir.AttrBuilder.get('BoolAttr')(unit_diagonal, context=_ods_context)) attributes["transpose_a"] = (transpose_a if ( isinstance(transpose_a, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('MHLO_TransposeAttr')) else _ods_ir.AttrBuilder.get('MHLO_TransposeAttr')(transpose_a, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def a(self): return self.operation.operands[0] @builtins.property def b(self): return self.operation.operands[1] @builtins.property def left_side(self): return self.operation.attributes["left_side"] @left_side.setter def left_side(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["left_side"] = value @builtins.property def lower(self): return self.operation.attributes["lower"] @lower.setter def lower(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["lower"] = value @builtins.property def unit_diagonal(self): return self.operation.attributes["unit_diagonal"] @unit_diagonal.setter def unit_diagonal(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["unit_diagonal"] = value @builtins.property def transpose_a(self): return self.operation.attributes["transpose_a"] @transpose_a.setter def transpose_a(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["transpose_a"] = value def triangular_solve(a, b, left_side, lower, unit_diagonal, transpose_a, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(TriangularSolveOp(a=a, b=b, left_side=left_side, lower=lower, unit_diagonal=unit_diagonal, transpose_a=transpose_a, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class TupleOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.tuple" _ODS_REGIONS = (0, True) def __init__(self, val, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(val)) _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 val(self): _ods_variadic_group_length = len(self.operation.operands) - 1 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @builtins.property def result(self): return self.operation.results[0] def tuple(val, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(TupleOp(val=val, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class UniformDequantizeOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.uniform_dequantize" _ODS_REGIONS = (0, True) def __init__(self, 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) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, 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] def uniform_dequantize(operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(UniformDequantizeOp(operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class UniformQuantizeOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.uniform_quantize" _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] def uniform_quantize(result, operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(UniformQuantizeOp(result=result, operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class WhileOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.while" _ODS_REGIONS = (2, True) def __init__(self, result, operand, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None operands.extend(_get_op_results_or_values(operand)) _ods_context = _ods_get_default_loc_context(loc) results.extend(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): _ods_variadic_group_length = len(self.operation.operands) - 1 + 1 return self.operation.operands[0:0 + _ods_variadic_group_length] @builtins.property def cond(self): return self.regions[0] @builtins.property def body(self): return self.regions[1] def while_(result, operand, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(WhileOp(result=result, operand=operand, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class XlaRngGetAndUpdateStateOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.xla.rng_get_and_update_state" _ODS_REGIONS = (0, True) def __init__(self, delta, *, loc=None, ip=None): operands = [] results = [] attributes = {} regions = None _ods_context = _ods_get_default_loc_context(loc) attributes["delta"] = (delta if ( isinstance(delta, _ods_ir.Attribute) or not _ods_ir.AttrBuilder.contains('I64Attr')) else _ods_ir.AttrBuilder.get('I64Attr')(delta, context=_ods_context)) _ods_successors = None super().__init__(self.build_generic(attributes=attributes, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)) @builtins.property def delta(self): return self.operation.attributes["delta"] @delta.setter def delta(self, value): if value is None: raise ValueError("'None' not allowed as value for mandatory attributes") self.operation.attributes["delta"] = value def xla_rng_get_and_update_state(delta, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(XlaRngGetAndUpdateStateOp(delta=delta, loc=loc, ip=ip)) @_ods_cext.register_operation(_Dialect) class XorOp(_ods_ir.OpView): OPERATION_NAME = "mhlo.xor" _ODS_REGIONS = (0, True) def __init__(self, lhs, rhs, *, 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)) _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 lhs(self): return self.operation.operands[0] @builtins.property def rhs(self): return self.operation.operands[1] @builtins.property def result(self): return self.operation.results[0] def xor(lhs, rhs, *, loc=None, ip=None) -> _ods_ir.Value: return _get_op_result_or_op_results(XorOp(lhs=lhs, rhs=rhs, loc=loc, ip=ip))