""" @generated by mypy-protobuf. Do not edit manually! isort:skip_file This proto file defines messages which represent the HLO module. This is a full fidelity serialization of the c++ HLO constructs. Many of the protos below are simple 1-to-1 serializations of the corresponding C++ classes, e.g., HloModule, HloComputation, and HloInstruction. FIELD NAMES ARE IMPORTANT Unlike most protos, you can't safely change the names of fields, even if you keep the numeric ids the same. This is because we sometimes serialize these protos as JSON, which includes the field names in the serialization. """ import builtins import collections.abc import sys import typing import google.protobuf.descriptor import google.protobuf.internal.containers import google.protobuf.internal.enum_type_wrapper import google.protobuf.message import tensorflow.compiler.xla.xla_data_pb2 if sys.version_info >= (3, 10): import typing as typing_extensions else: import typing_extensions DESCRIPTOR: google.protobuf.descriptor.FileDescriptor class _CustomCallSchedule: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _CustomCallScheduleEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_CustomCallSchedule.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor SCHEDULE_NONE: _CustomCallSchedule.ValueType # 0 SCHEDULE_LATEST: _CustomCallSchedule.ValueType # 1 SCHEDULE_EARLIEST: _CustomCallSchedule.ValueType # 2 class CustomCallSchedule(_CustomCallSchedule, metaclass=_CustomCallScheduleEnumTypeWrapper): ... SCHEDULE_NONE: CustomCallSchedule.ValueType # 0 SCHEDULE_LATEST: CustomCallSchedule.ValueType # 1 SCHEDULE_EARLIEST: CustomCallSchedule.ValueType # 2 global___CustomCallSchedule = CustomCallSchedule class _CustomCallApiVersion: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _CustomCallApiVersionEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_CustomCallApiVersion.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor API_VERSION_UNSPECIFIED: _CustomCallApiVersion.ValueType # 0 API_VERSION_ORIGINAL: _CustomCallApiVersion.ValueType # 1 """The first version of the API, with the following signatures: CPU: void do_custom_call(void* out, const void** in); GPU: void do_custom_call(CUstream stream, void** buffers, const char* opaque, size_t opaque_len); """ API_VERSION_STATUS_RETURNING: _CustomCallApiVersion.ValueType # 2 """When the ability to return success/failure status was added: CPU: void do_custom_call(void* out, const void** in, XlaCustomCallStatus* status); GPU: void do_custom_call(CUstream stream, void** buffers, const char* opaque, size_t opaque_len, XlaCustomCallStatus* status); """ API_VERSION_STATUS_RETURNING_UNIFIED: _CustomCallApiVersion.ValueType # 3 """Fixes the API signatures on the CPU side of the version STATUS_RETURNING by adding the opaque string so that the custom call API is consistent across CPUs and GPUs. For GPUs, the behaviors invoked by API_VERSION_STATUS_RETURNING and API_VERSION_STATUS_RETURNING_UNIFIED are the same. CPU: void do_custom_call(void* out, const void** in, const char* opaque, size_t opaque_len, XlaCustomCallStatus* status); GPU: void do_custom_call(CUstream stream, void** buffers, const char* opaque, size_t opaque_len, XlaCustomCallStatus* status); """ API_VERSION_TYPED_FFI: _CustomCallApiVersion.ValueType # 4 """Api version implementing XLA runtime custom call calling convention. These custom calls can be registered as an XLA runtime custom call (1) or as XLA runtime FFI binding (2). This type of custom call uses custom ABI to pass type information along with custom call arguments. Also it passes buffer arguments together with data type, sizes and strides. Example: (XLA runtime custom call) absl::Status DoCustomCall(StridedMemrefView arg, float attr); CustomCall::Bind("custom_call") .Arg() .Attr("attr") .To(DoCustomCall); (1) xla/runtime/custom_call.h (2) xla/runtime/ffi/ffi.h """ class CustomCallApiVersion(_CustomCallApiVersion, metaclass=_CustomCallApiVersionEnumTypeWrapper): """The version of the API used by the custom call function. The signatures for each version are given below. TODO(b/189822916): Remove this enum when all clients are migrated to the status-returning API. """ API_VERSION_UNSPECIFIED: CustomCallApiVersion.ValueType # 0 API_VERSION_ORIGINAL: CustomCallApiVersion.ValueType # 1 """The first version of the API, with the following signatures: CPU: void do_custom_call(void* out, const void** in); GPU: void do_custom_call(CUstream stream, void** buffers, const char* opaque, size_t opaque_len); """ API_VERSION_STATUS_RETURNING: CustomCallApiVersion.ValueType # 2 """When the ability to return success/failure status was added: CPU: void do_custom_call(void* out, const void** in, XlaCustomCallStatus* status); GPU: void do_custom_call(CUstream stream, void** buffers, const char* opaque, size_t opaque_len, XlaCustomCallStatus* status); """ API_VERSION_STATUS_RETURNING_UNIFIED: CustomCallApiVersion.ValueType # 3 """Fixes the API signatures on the CPU side of the version STATUS_RETURNING by adding the opaque string so that the custom call API is consistent across CPUs and GPUs. For GPUs, the behaviors invoked by API_VERSION_STATUS_RETURNING and API_VERSION_STATUS_RETURNING_UNIFIED are the same. CPU: void do_custom_call(void* out, const void** in, const char* opaque, size_t opaque_len, XlaCustomCallStatus* status); GPU: void do_custom_call(CUstream stream, void** buffers, const char* opaque, size_t opaque_len, XlaCustomCallStatus* status); """ API_VERSION_TYPED_FFI: CustomCallApiVersion.ValueType # 4 """Api version implementing XLA runtime custom call calling convention. These custom calls can be registered as an XLA runtime custom call (1) or as XLA runtime FFI binding (2). This type of custom call uses custom ABI to pass type information along with custom call arguments. Also it passes buffer arguments together with data type, sizes and strides. Example: (XLA runtime custom call) absl::Status DoCustomCall(StridedMemrefView arg, float attr); CustomCall::Bind("custom_call") .Arg() .Attr("attr") .To(DoCustomCall); (1) xla/runtime/custom_call.h (2) xla/runtime/ffi/ffi.h """ global___CustomCallApiVersion = CustomCallApiVersion class _Kind: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _KindEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_Kind.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor UNDEFINED_ALIAS: _Kind.ValueType # 0 """Define a UNDEFINED_ALIAS equal to zero to get around the default-0 proto3 behavior and missing has_*() APIs. """ MAY_ALIAS: _Kind.ValueType # 1 """The buffers may or may not alias at runtime.""" MUST_ALIAS: _Kind.ValueType # 2 """The buffers must alias at runtime.""" class Kind(_Kind, metaclass=_KindEnumTypeWrapper): ... UNDEFINED_ALIAS: Kind.ValueType # 0 """Define a UNDEFINED_ALIAS equal to zero to get around the default-0 proto3 behavior and missing has_*() APIs. """ MAY_ALIAS: Kind.ValueType # 1 """The buffers may or may not alias at runtime.""" MUST_ALIAS: Kind.ValueType # 2 """The buffers must alias at runtime.""" global___Kind = Kind @typing_extensions.final class HloInstructionProto(google.protobuf.message.Message): """Serialization of HloInstruction. Next ID: 81 """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class SliceDimensions(google.protobuf.message.Message): """Describes the [begin, end) index range and stride for slices.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor START_FIELD_NUMBER: builtins.int LIMIT_FIELD_NUMBER: builtins.int STRIDE_FIELD_NUMBER: builtins.int start: builtins.int limit: builtins.int stride: builtins.int def __init__( self, *, start: builtins.int | None = ..., limit: builtins.int | None = ..., stride: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["limit", b"limit", "start", b"start", "stride", b"stride"]) -> None: ... NAME_FIELD_NUMBER: builtins.int OPCODE_FIELD_NUMBER: builtins.int SHAPE_FIELD_NUMBER: builtins.int METADATA_FIELD_NUMBER: builtins.int LITERAL_FIELD_NUMBER: builtins.int PARAMETER_NUMBER_FIELD_NUMBER: builtins.int FUSION_KIND_FIELD_NUMBER: builtins.int TUPLE_INDEX_FIELD_NUMBER: builtins.int DIMENSIONS_FIELD_NUMBER: builtins.int WINDOW_FIELD_NUMBER: builtins.int CONVOLUTION_DIMENSION_NUMBERS_FIELD_NUMBER: builtins.int FEATURE_GROUP_COUNT_FIELD_NUMBER: builtins.int BATCH_GROUP_COUNT_FIELD_NUMBER: builtins.int SLICE_DIMENSIONS_FIELD_NUMBER: builtins.int EXPONENT_BITS_FIELD_NUMBER: builtins.int MANTISSA_BITS_FIELD_NUMBER: builtins.int DYNAMIC_SLICE_SIZES_FIELD_NUMBER: builtins.int PADDING_CONFIG_FIELD_NUMBER: builtins.int OUTFEED_CONFIG_FIELD_NUMBER: builtins.int DISTRIBUTION_FIELD_NUMBER: builtins.int EPSILON_FIELD_NUMBER: builtins.int FEATURE_INDEX_FIELD_NUMBER: builtins.int CHANNEL_ID_FIELD_NUMBER: builtins.int INFEED_CONFIG_FIELD_NUMBER: builtins.int CUSTOM_CALL_TARGET_FIELD_NUMBER: builtins.int OUTFEED_SHAPE_FIELD_NUMBER: builtins.int DOT_DIMENSION_NUMBERS_FIELD_NUMBER: builtins.int FFT_TYPE_FIELD_NUMBER: builtins.int FFT_LENGTH_FIELD_NUMBER: builtins.int COMPARISON_DIRECTION_FIELD_NUMBER: builtins.int GATHER_DIMENSION_NUMBERS_FIELD_NUMBER: builtins.int GATHER_SLICE_SIZES_FIELD_NUMBER: builtins.int ID_FIELD_NUMBER: builtins.int OPERAND_IDS_FIELD_NUMBER: builtins.int CONTROL_PREDECESSOR_IDS_FIELD_NUMBER: builtins.int CALLED_COMPUTATION_IDS_FIELD_NUMBER: builtins.int SHARDING_FIELD_NUMBER: builtins.int BACKEND_CONFIG_FIELD_NUMBER: builtins.int REPLICA_GROUPS_FIELD_NUMBER: builtins.int ALL_REDUCE_ID_FIELD_NUMBER: builtins.int USE_GLOBAL_DEVICE_IDS_FIELD_NUMBER: builtins.int IS_HOST_TRANSFER_FIELD_NUMBER: builtins.int IS_STABLE_FIELD_NUMBER: builtins.int SCATTER_DIMENSION_NUMBERS_FIELD_NUMBER: builtins.int PRECISION_CONFIG_FIELD_NUMBER: builtins.int SOURCE_TARGET_PAIRS_FIELD_NUMBER: builtins.int DOMAIN_ENTRY_SHARDING_FIELD_NUMBER: builtins.int DOMAIN_EXIT_SHARDING_FIELD_NUMBER: builtins.int CONSTRAIN_LAYOUT_FIELD_NUMBER: builtins.int OPERAND_SHAPES_WITH_LAYOUT_FIELD_NUMBER: builtins.int TRIANGULAR_SOLVE_OPTIONS_FIELD_NUMBER: builtins.int CHOLESKY_OPTIONS_FIELD_NUMBER: builtins.int PARAMETER_REPLICATION_FIELD_NUMBER: builtins.int CUSTOM_CALL_HAS_SIDE_EFFECT_FIELD_NUMBER: builtins.int OUTPUT_OPERAND_ALIASING_FIELD_NUMBER: builtins.int CUSTOM_CALL_SCHEDULE_FIELD_NUMBER: builtins.int DELTA_FIELD_NUMBER: builtins.int INDICES_ARE_SORTED_FIELD_NUMBER: builtins.int FRONTEND_ATTRIBUTES_FIELD_NUMBER: builtins.int UNIQUE_INDICES_FIELD_NUMBER: builtins.int RNG_ALGORITHM_FIELD_NUMBER: builtins.int COMPARISON_TYPE_FIELD_NUMBER: builtins.int IS_CROSS_PROGRAM_PREFETCH_FIELD_NUMBER: builtins.int CROSS_PROGRAM_PREFETCH_INDEX_FIELD_NUMBER: builtins.int PADDING_TYPE_FIELD_NUMBER: builtins.int CUSTOM_CALL_API_VERSION_FIELD_NUMBER: builtins.int ASYNC_GROUP_ID_FIELD_NUMBER: builtins.int ASYNC_EXECUTION_THREAD_FIELD_NUMBER: builtins.int name: builtins.str opcode: builtins.str @property def shape(self) -> tensorflow.compiler.xla.xla_data_pb2.ShapeProto: ... @property def metadata(self) -> tensorflow.compiler.xla.xla_data_pb2.OpMetadata: ... @property def literal(self) -> tensorflow.compiler.xla.xla_data_pb2.LiteralProto: """Literal, only present for kConstant.""" parameter_number: builtins.int """Parameter number is only present for kParameter.""" fusion_kind: builtins.str """Fusion state, only present for kFusion.""" tuple_index: builtins.int """Index for kGetTupleElement.""" @property def dimensions(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: """Dimensions present for some operations that require reshaping or broadcasting, including Reshape, Reduce, ReduceWindow, and Reverse. """ @property def window(self) -> tensorflow.compiler.xla.xla_data_pb2.Window: """Describes the window in a windowed operation such as convolution.""" @property def convolution_dimension_numbers(self) -> tensorflow.compiler.xla.xla_data_pb2.ConvolutionDimensionNumbers: """Describes the dimension numbers used for a convolution.""" feature_group_count: builtins.int """The number of feature groups. Used for a convolution. Must be a divisor of the input feature dimension and output feature dimension. If not specified, it will use a default value of 1. """ batch_group_count: builtins.int @property def slice_dimensions(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___HloInstructionProto.SliceDimensions]: ... exponent_bits: builtins.int """The bit sizes for a reduce-precision operation.""" mantissa_bits: builtins.int @property def dynamic_slice_sizes(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: """Describes the [start, start + size) range size for a dynamic slice ('start' is specified dynamically in the second operand of the operation). """ @property def padding_config(self) -> tensorflow.compiler.xla.xla_data_pb2.PaddingConfig: """The padding configuration that describes the edge padding and interior padding of this pad instruction. Only set for pad instructions. """ outfeed_config: builtins.bytes """Outfeed configuration information, only present for kOutfeed.""" distribution: tensorflow.compiler.xla.xla_data_pb2.RandomDistribution.ValueType """The distribution requested for random number generation. Only present for kRng. """ epsilon: builtins.float """A small float number added to the variance to avoid divide-by-zero error. Only present for kBatchNormTraining, kBatchNormInference, and kBatchNormGrad. """ feature_index: builtins.int """An integer value representing the index of the feature dimension. Only present for kBatchNormTraining, kBatchNormInference, and kBatchNormGrad. """ channel_id: builtins.int """Represents a unique identifier for each Send/Recv instruction pair or optionally for collective instructions (AllReduce, CollectivePermute, AllToAll). Non-positive channel_id is equivalent to no channel id. """ infeed_config: builtins.bytes """The string representation of the infeed configuration.""" custom_call_target: builtins.str """Name of a external target (eg, global symbol) to call, only present for kCustomCall. """ @property def outfeed_shape(self) -> tensorflow.compiler.xla.xla_data_pb2.ShapeProto: """Shape of outfeed request.""" @property def dot_dimension_numbers(self) -> tensorflow.compiler.xla.xla_data_pb2.DotDimensionNumbers: """Describes the dimension numbers used for a dot operation""" fft_type: tensorflow.compiler.xla.xla_data_pb2.FftType.ValueType """FFT type (FFT, IFFT, etc).""" @property def fft_length(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: """FFT length.""" comparison_direction: builtins.str """Comparison direction only used for kCompare.""" @property def gather_dimension_numbers(self) -> tensorflow.compiler.xla.xla_data_pb2.GatherDimensionNumbers: """Gather dimension numbers.""" @property def gather_slice_sizes(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ... id: builtins.int """The id of this instruction.""" @property def operand_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ... @property def control_predecessor_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ... @property def called_computation_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ... @property def sharding(self) -> tensorflow.compiler.xla.xla_data_pb2.OpSharding: ... backend_config: builtins.bytes """Backend configuration for the instruction. Has backend-specific meaning.""" @property def replica_groups(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[tensorflow.compiler.xla.xla_data_pb2.ReplicaGroup]: """Cross replica op fields.""" all_reduce_id: builtins.int """Deprecated, but keeping it for backward compatibility. Use channel_id. Non-positive all_reduce_id is equivalent to no all_reduce_id. """ use_global_device_ids: builtins.bool """If true, interprets ids in ReplicaGroup as global device ids, which is a linearized id of `replica_id * partition_count + partition_id`. """ is_host_transfer: builtins.bool """Whether this Send/Recv instruction transfers data to/from the host. Only present for Send and Recv instructions and their SendDone and RecvDone partners. """ is_stable: builtins.bool """Whether this Sort instruction should be stable.""" @property def scatter_dimension_numbers(self) -> tensorflow.compiler.xla.xla_data_pb2.ScatterDimensionNumbers: ... @property def precision_config(self) -> tensorflow.compiler.xla.xla_data_pb2.PrecisionConfig: """Precision configuration for the instruction. Has backend-specific meaning.""" @property def source_target_pairs(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[tensorflow.compiler.xla.xla_data_pb2.SourceTarget]: """Collective permute field.""" @property def domain_entry_sharding(self) -> tensorflow.compiler.xla.xla_data_pb2.OpSharding: """Sharding for kDomain instructions.""" @property def domain_exit_sharding(self) -> tensorflow.compiler.xla.xla_data_pb2.OpSharding: ... constrain_layout: builtins.bool """For custom call this indicates that the layouts are constrained. If constrain_layout is true then the 'shape' field must contain a layout, and 'operand_shapes_with_layout' must contain a shape with layout for each operand. """ @property def operand_shapes_with_layout(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[tensorflow.compiler.xla.xla_data_pb2.ShapeProto]: ... @property def triangular_solve_options(self) -> tensorflow.compiler.xla.xla_data_pb2.TriangularSolveOptions: """Options for TriangularSolve""" @property def cholesky_options(self) -> tensorflow.compiler.xla.xla_data_pb2.CholeskyOptions: """Options for Cholesky""" @property def parameter_replication(self) -> tensorflow.compiler.xla.xla_data_pb2.ParameterReplication: """Describes how parameters behave with regards to replicas.""" custom_call_has_side_effect: builtins.bool """Whether the kCustomCall instruction has side-effects, only present for kCustomCall. """ @property def output_operand_aliasing(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[tensorflow.compiler.xla.xla_data_pb2.OutputOperandAliasing]: """A list of OutputOperandAliasing pairs that specifies aliasing buffers between output and operands for kCustomCall and kFusion. """ custom_call_schedule: global___CustomCallSchedule.ValueType """Specifies the desired schedule for the custom-call. The field is only present for custom-call. """ delta: builtins.int """The delta value for kRngGetAndUpdateState.""" indices_are_sorted: builtins.bool """Specifies if the gather/scatter indices are guaranteed to be sorted by the caller. """ @property def frontend_attributes(self) -> tensorflow.compiler.xla.xla_data_pb2.FrontendAttributes: """Frontend attributes to pass to the XLA backend.""" unique_indices: builtins.bool """Specifies if all elements updated are guaranteed to be unique by the caller. """ rng_algorithm: tensorflow.compiler.xla.xla_data_pb2.RandomAlgorithm.ValueType """RNG algorithm used by kRngBitGenerator.""" comparison_type: builtins.str """The comparison type used for kCompare.""" is_cross_program_prefetch: builtins.bool """Specifies if this is a cross-program-prefetch, used by kCopyStart. Deprecated and replaced by optional_cross_program_prefetch_index. """ cross_program_prefetch_index: builtins.int padding_type: tensorflow.compiler.xla.xla_data_pb2.PaddingType.ValueType """If a convolution is dynamic, a dynamic padding type will be specified.""" custom_call_api_version: global___CustomCallApiVersion.ValueType """The API version used by the custom call function. This field is only present for custom-call. TODO(b/189822916): Remove this field when all clients are migrated to the status-returning API. """ async_group_id: builtins.int """Represents a unique identifier for an async group which consists of an async start, async done, and zero or more async update operations. Negative async_group_id is equivalent to no async group id. """ async_execution_thread: builtins.str """Represents a unique execution thread name for one or more async groups. Each HLO module may contain a main thread and one or more parallel threads. Empty async_execution_thread is equivalent to main thread. """ def __init__( self, *, name: builtins.str | None = ..., opcode: builtins.str | None = ..., shape: tensorflow.compiler.xla.xla_data_pb2.ShapeProto | None = ..., metadata: tensorflow.compiler.xla.xla_data_pb2.OpMetadata | None = ..., literal: tensorflow.compiler.xla.xla_data_pb2.LiteralProto | None = ..., parameter_number: builtins.int | None = ..., fusion_kind: builtins.str | None = ..., tuple_index: builtins.int | None = ..., dimensions: collections.abc.Iterable[builtins.int] | None = ..., window: tensorflow.compiler.xla.xla_data_pb2.Window | None = ..., convolution_dimension_numbers: tensorflow.compiler.xla.xla_data_pb2.ConvolutionDimensionNumbers | None = ..., feature_group_count: builtins.int | None = ..., batch_group_count: builtins.int | None = ..., slice_dimensions: collections.abc.Iterable[global___HloInstructionProto.SliceDimensions] | None = ..., exponent_bits: builtins.int | None = ..., mantissa_bits: builtins.int | None = ..., dynamic_slice_sizes: collections.abc.Iterable[builtins.int] | None = ..., padding_config: tensorflow.compiler.xla.xla_data_pb2.PaddingConfig | None = ..., outfeed_config: builtins.bytes | None = ..., distribution: tensorflow.compiler.xla.xla_data_pb2.RandomDistribution.ValueType | None = ..., epsilon: builtins.float | None = ..., feature_index: builtins.int | None = ..., channel_id: builtins.int | None = ..., infeed_config: builtins.bytes | None = ..., custom_call_target: builtins.str | None = ..., outfeed_shape: tensorflow.compiler.xla.xla_data_pb2.ShapeProto | None = ..., dot_dimension_numbers: tensorflow.compiler.xla.xla_data_pb2.DotDimensionNumbers | None = ..., fft_type: tensorflow.compiler.xla.xla_data_pb2.FftType.ValueType | None = ..., fft_length: collections.abc.Iterable[builtins.int] | None = ..., comparison_direction: builtins.str | None = ..., gather_dimension_numbers: tensorflow.compiler.xla.xla_data_pb2.GatherDimensionNumbers | None = ..., gather_slice_sizes: collections.abc.Iterable[builtins.int] | None = ..., id: builtins.int | None = ..., operand_ids: collections.abc.Iterable[builtins.int] | None = ..., control_predecessor_ids: collections.abc.Iterable[builtins.int] | None = ..., called_computation_ids: collections.abc.Iterable[builtins.int] | None = ..., sharding: tensorflow.compiler.xla.xla_data_pb2.OpSharding | None = ..., backend_config: builtins.bytes | None = ..., replica_groups: collections.abc.Iterable[tensorflow.compiler.xla.xla_data_pb2.ReplicaGroup] | None = ..., all_reduce_id: builtins.int | None = ..., use_global_device_ids: builtins.bool | None = ..., is_host_transfer: builtins.bool | None = ..., is_stable: builtins.bool | None = ..., scatter_dimension_numbers: tensorflow.compiler.xla.xla_data_pb2.ScatterDimensionNumbers | None = ..., precision_config: tensorflow.compiler.xla.xla_data_pb2.PrecisionConfig | None = ..., source_target_pairs: collections.abc.Iterable[tensorflow.compiler.xla.xla_data_pb2.SourceTarget] | None = ..., domain_entry_sharding: tensorflow.compiler.xla.xla_data_pb2.OpSharding | None = ..., domain_exit_sharding: tensorflow.compiler.xla.xla_data_pb2.OpSharding | None = ..., constrain_layout: builtins.bool | None = ..., operand_shapes_with_layout: collections.abc.Iterable[tensorflow.compiler.xla.xla_data_pb2.ShapeProto] | None = ..., triangular_solve_options: tensorflow.compiler.xla.xla_data_pb2.TriangularSolveOptions | None = ..., cholesky_options: tensorflow.compiler.xla.xla_data_pb2.CholeskyOptions | None = ..., parameter_replication: tensorflow.compiler.xla.xla_data_pb2.ParameterReplication | None = ..., custom_call_has_side_effect: builtins.bool | None = ..., output_operand_aliasing: collections.abc.Iterable[tensorflow.compiler.xla.xla_data_pb2.OutputOperandAliasing] | None = ..., custom_call_schedule: global___CustomCallSchedule.ValueType | None = ..., delta: builtins.int | None = ..., indices_are_sorted: builtins.bool | None = ..., frontend_attributes: tensorflow.compiler.xla.xla_data_pb2.FrontendAttributes | None = ..., unique_indices: builtins.bool | None = ..., rng_algorithm: tensorflow.compiler.xla.xla_data_pb2.RandomAlgorithm.ValueType | None = ..., comparison_type: builtins.str | None = ..., is_cross_program_prefetch: builtins.bool | None = ..., cross_program_prefetch_index: builtins.int | None = ..., padding_type: tensorflow.compiler.xla.xla_data_pb2.PaddingType.ValueType | None = ..., custom_call_api_version: global___CustomCallApiVersion.ValueType | None = ..., async_group_id: builtins.int | None = ..., async_execution_thread: builtins.str | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["cholesky_options", b"cholesky_options", "convolution_dimension_numbers", b"convolution_dimension_numbers", "cross_program_prefetch_index", b"cross_program_prefetch_index", "domain_entry_sharding", b"domain_entry_sharding", "domain_exit_sharding", b"domain_exit_sharding", "dot_dimension_numbers", b"dot_dimension_numbers", "frontend_attributes", b"frontend_attributes", "gather_dimension_numbers", b"gather_dimension_numbers", "literal", b"literal", "metadata", b"metadata", "optional_cross_program_prefetch_index", b"optional_cross_program_prefetch_index", "outfeed_shape", b"outfeed_shape", "padding_config", b"padding_config", "parameter_replication", b"parameter_replication", "precision_config", b"precision_config", "scatter_dimension_numbers", b"scatter_dimension_numbers", "shape", b"shape", "sharding", b"sharding", "triangular_solve_options", b"triangular_solve_options", "window", b"window"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["all_reduce_id", b"all_reduce_id", "async_execution_thread", b"async_execution_thread", "async_group_id", b"async_group_id", "backend_config", b"backend_config", "batch_group_count", b"batch_group_count", "called_computation_ids", b"called_computation_ids", "channel_id", b"channel_id", "cholesky_options", b"cholesky_options", "comparison_direction", b"comparison_direction", "comparison_type", b"comparison_type", "constrain_layout", b"constrain_layout", "control_predecessor_ids", b"control_predecessor_ids", "convolution_dimension_numbers", b"convolution_dimension_numbers", "cross_program_prefetch_index", b"cross_program_prefetch_index", "custom_call_api_version", b"custom_call_api_version", "custom_call_has_side_effect", b"custom_call_has_side_effect", "custom_call_schedule", b"custom_call_schedule", "custom_call_target", b"custom_call_target", "delta", b"delta", "dimensions", b"dimensions", "distribution", b"distribution", "domain_entry_sharding", b"domain_entry_sharding", "domain_exit_sharding", b"domain_exit_sharding", "dot_dimension_numbers", b"dot_dimension_numbers", "dynamic_slice_sizes", b"dynamic_slice_sizes", "epsilon", b"epsilon", "exponent_bits", b"exponent_bits", "feature_group_count", b"feature_group_count", "feature_index", b"feature_index", "fft_length", b"fft_length", "fft_type", b"fft_type", "frontend_attributes", b"frontend_attributes", "fusion_kind", b"fusion_kind", "gather_dimension_numbers", b"gather_dimension_numbers", "gather_slice_sizes", b"gather_slice_sizes", "id", b"id", "indices_are_sorted", b"indices_are_sorted", "infeed_config", b"infeed_config", "is_cross_program_prefetch", b"is_cross_program_prefetch", "is_host_transfer", b"is_host_transfer", "is_stable", b"is_stable", "literal", b"literal", "mantissa_bits", b"mantissa_bits", "metadata", b"metadata", "name", b"name", "opcode", b"opcode", "operand_ids", b"operand_ids", "operand_shapes_with_layout", b"operand_shapes_with_layout", "optional_cross_program_prefetch_index", b"optional_cross_program_prefetch_index", "outfeed_config", b"outfeed_config", "outfeed_shape", b"outfeed_shape", "output_operand_aliasing", b"output_operand_aliasing", "padding_config", b"padding_config", "padding_type", b"padding_type", "parameter_number", b"parameter_number", "parameter_replication", b"parameter_replication", "precision_config", b"precision_config", "replica_groups", b"replica_groups", "rng_algorithm", b"rng_algorithm", "scatter_dimension_numbers", b"scatter_dimension_numbers", "shape", b"shape", "sharding", b"sharding", "slice_dimensions", b"slice_dimensions", "source_target_pairs", b"source_target_pairs", "triangular_solve_options", b"triangular_solve_options", "tuple_index", b"tuple_index", "unique_indices", b"unique_indices", "use_global_device_ids", b"use_global_device_ids", "window", b"window"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["optional_cross_program_prefetch_index", b"optional_cross_program_prefetch_index"]) -> typing_extensions.Literal["cross_program_prefetch_index"] | None: ... global___HloInstructionProto = HloInstructionProto @typing_extensions.final class HloComputationProto(google.protobuf.message.Message): """Serialization of HloComputation.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor NAME_FIELD_NUMBER: builtins.int INSTRUCTIONS_FIELD_NUMBER: builtins.int PROGRAM_SHAPE_FIELD_NUMBER: builtins.int ID_FIELD_NUMBER: builtins.int ROOT_ID_FIELD_NUMBER: builtins.int IS_FUSION_COMPUTATION_FIELD_NUMBER: builtins.int EXECUTION_THREAD_FIELD_NUMBER: builtins.int name: builtins.str @property def instructions(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___HloInstructionProto]: """The array of instructions is always in a valid dependency order, where operands appear before their users. """ @property def program_shape(self) -> tensorflow.compiler.xla.xla_data_pb2.ProgramShapeProto: """The program shape (with layout) of this computation.""" id: builtins.int """The id of this computation.""" root_id: builtins.int """The id of the root of the computation.""" is_fusion_computation: builtins.bool """Whether this is a fusion computation. Fusion computations should use this to determine whether they are a fusion in CreateFromProto since the parent fusion_instruction_ may get removed and be nullptr. """ execution_thread: builtins.str """The name of execution thread this computation belongs to.""" def __init__( self, *, name: builtins.str | None = ..., instructions: collections.abc.Iterable[global___HloInstructionProto] | None = ..., program_shape: tensorflow.compiler.xla.xla_data_pb2.ProgramShapeProto | None = ..., id: builtins.int | None = ..., root_id: builtins.int | None = ..., is_fusion_computation: builtins.bool | None = ..., execution_thread: builtins.str | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["program_shape", b"program_shape"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["execution_thread", b"execution_thread", "id", b"id", "instructions", b"instructions", "is_fusion_computation", b"is_fusion_computation", "name", b"name", "program_shape", b"program_shape", "root_id", b"root_id"]) -> None: ... global___HloComputationProto = HloComputationProto @typing_extensions.final class HloScheduleProto(google.protobuf.message.Message): """Serialization of an HLO schedule. An HLO schedule contains a total order of instructions for each non-fusion computation in the module. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class InstructionSequence(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor INSTRUCTION_IDS_FIELD_NUMBER: builtins.int @property def instruction_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ... def __init__( self, *, instruction_ids: collections.abc.Iterable[builtins.int] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["instruction_ids", b"instruction_ids"]) -> None: ... @typing_extensions.final class SequencesEntry(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor KEY_FIELD_NUMBER: builtins.int VALUE_FIELD_NUMBER: builtins.int key: builtins.int @property def value(self) -> global___HloScheduleProto.InstructionSequence: ... def __init__( self, *, key: builtins.int | None = ..., value: global___HloScheduleProto.InstructionSequence | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... SEQUENCES_FIELD_NUMBER: builtins.int @property def sequences(self) -> google.protobuf.internal.containers.MessageMap[builtins.int, global___HloScheduleProto.InstructionSequence]: """Map from computation id to sequence.""" def __init__( self, *, sequences: collections.abc.Mapping[builtins.int, global___HloScheduleProto.InstructionSequence] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["sequences", b"sequences"]) -> None: ... global___HloScheduleProto = HloScheduleProto @typing_extensions.final class HloInputOutputAliasProto(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class AliasEntryProto(google.protobuf.message.Message): """The following proto describes a pair of aliased an input (described by parameter number and a ShapeIndex of the parameter) and an output (described by a ShapeIndex of the root instruction). For example: entry = { output_shape_index={1}, parameter_number=0, parameter_shape_index={1, 2}, } This entry indicates that the first paremter's {1, 2} element is aliased with the {1} element of the root instruction. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor OUTPUT_SHAPE_INDEX_FIELD_NUMBER: builtins.int PARAMETER_NUMBER_FIELD_NUMBER: builtins.int PARAMETER_SHAPE_INDEX_FIELD_NUMBER: builtins.int KIND_FIELD_NUMBER: builtins.int @property def output_shape_index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: """ShapeIndex of the root hlo.""" parameter_number: builtins.int """Number of the parameter in entry computation.""" @property def parameter_shape_index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: """ShapeIndex of the parameter instruction.""" kind: global___Kind.ValueType """The kind of alias to be setup.""" def __init__( self, *, output_shape_index: collections.abc.Iterable[builtins.int] | None = ..., parameter_number: builtins.int | None = ..., parameter_shape_index: collections.abc.Iterable[builtins.int] | None = ..., kind: global___Kind.ValueType | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["kind", b"kind", "output_shape_index", b"output_shape_index", "parameter_number", b"parameter_number", "parameter_shape_index", b"parameter_shape_index"]) -> None: ... ENTRIES_FIELD_NUMBER: builtins.int @property def entries(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___HloInputOutputAliasProto.AliasEntryProto]: ... def __init__( self, *, entries: collections.abc.Iterable[global___HloInputOutputAliasProto.AliasEntryProto] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["entries", b"entries"]) -> None: ... global___HloInputOutputAliasProto = HloInputOutputAliasProto @typing_extensions.final class DynamicParameterBindingProto(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class Binding(google.protobuf.message.Message): """A list of bindings which indicates that the `target_param_dim_num` in the subshape `target_param_index` of parameter `target_param_num` is a dynamic dimension and its real dynamic size is represented by `dynamic_param_index` in parameter `dynamic_param_num`. As an example, imagine we have a program: ENTRY main { a = f32[] parameter(0) b = f32[10] parameter(1) ROOT root = (f32[], f32[10]) tuple(%a, %b) } Let's say 'b' (param index 1) is a dynamic shape whose input has an upperbound of 10 and real size is determined at runtime.'a' represents the real size of b's first dimension. In this case, the fields are set in the following way: dynamic_param_num = 1 dynamic_param_index = {} target_param_num = 0 target_param_index = {} target_param_dim_num = 0 """ DESCRIPTOR: google.protobuf.descriptor.Descriptor DYNAMIC_PARAM_NUM_FIELD_NUMBER: builtins.int DYNAMIC_PARAM_INDEX_FIELD_NUMBER: builtins.int TARGET_PARAM_NUM_FIELD_NUMBER: builtins.int TARGET_PARAM_INDEX_FIELD_NUMBER: builtins.int TARGET_PARAM_DIM_NUM_FIELD_NUMBER: builtins.int dynamic_param_num: builtins.int @property def dynamic_param_index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ... target_param_num: builtins.int @property def target_param_index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ... target_param_dim_num: builtins.int def __init__( self, *, dynamic_param_num: builtins.int | None = ..., dynamic_param_index: collections.abc.Iterable[builtins.int] | None = ..., target_param_num: builtins.int | None = ..., target_param_index: collections.abc.Iterable[builtins.int] | None = ..., target_param_dim_num: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["dynamic_param_index", b"dynamic_param_index", "dynamic_param_num", b"dynamic_param_num", "target_param_dim_num", b"target_param_dim_num", "target_param_index", b"target_param_index", "target_param_num", b"target_param_num"]) -> None: ... ENTRIES_FIELD_NUMBER: builtins.int @property def entries(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___DynamicParameterBindingProto.Binding]: ... def __init__( self, *, entries: collections.abc.Iterable[global___DynamicParameterBindingProto.Binding] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["entries", b"entries"]) -> None: ... global___DynamicParameterBindingProto = DynamicParameterBindingProto @typing_extensions.final class CrossProgramPrefetch(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor PARAMETER_FIELD_NUMBER: builtins.int INDEX_FIELD_NUMBER: builtins.int OFFSET_FIELD_NUMBER: builtins.int parameter: builtins.int @property def index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ... offset: builtins.int def __init__( self, *, parameter: builtins.int | None = ..., index: collections.abc.Iterable[builtins.int] | None = ..., offset: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["index", b"index", "offset", b"offset", "parameter", b"parameter"]) -> None: ... global___CrossProgramPrefetch = CrossProgramPrefetch @typing_extensions.final class HloModuleProto(google.protobuf.message.Message): """Serialization of HloModule.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor class _ProfileType: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _ProfileTypeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[HloModuleProto._ProfileType.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor INVALID: HloModuleProto._ProfileType.ValueType # 0 FLAG: HloModuleProto._ProfileType.ValueType # 1 FUSION: HloModuleProto._ProfileType.ValueType # 2 LAYOUT: HloModuleProto._ProfileType.ValueType # 3 DOT: HloModuleProto._ProfileType.ValueType # 4 class ProfileType(_ProfileType, metaclass=_ProfileTypeEnumTypeWrapper): """The type of optimization profile in use for module-level optimizations.""" INVALID: HloModuleProto.ProfileType.ValueType # 0 FLAG: HloModuleProto.ProfileType.ValueType # 1 FUSION: HloModuleProto.ProfileType.ValueType # 2 LAYOUT: HloModuleProto.ProfileType.ValueType # 3 DOT: HloModuleProto.ProfileType.ValueType # 4 @typing_extensions.final class ProfileInfo(google.protobuf.message.Message): """Information about the optimization profile that this module contains.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor PROFILE_TYPE_FIELD_NUMBER: builtins.int RELATIVE_SPEEDUP_FIELD_NUMBER: builtins.int PROFILE_SOURCE_FIELD_NUMBER: builtins.int COMPILATION_EVENT_FIELD_NUMBER: builtins.int profile_type: global___HloModuleProto.ProfileType.ValueType """The optimization profiles that this module contains.""" relative_speedup: builtins.float """Speedup of tuned config compared to default config.""" profile_source: tensorflow.compiler.xla.xla_data_pb2.ProfileSource.ValueType """The source of the optimization profile that this module contains.""" compilation_event: tensorflow.compiler.xla.xla_data_pb2.CompilationEvent.ValueType """The compilation event that triggered the use of the profile.""" def __init__( self, *, profile_type: global___HloModuleProto.ProfileType.ValueType | None = ..., relative_speedup: builtins.float | None = ..., profile_source: tensorflow.compiler.xla.xla_data_pb2.ProfileSource.ValueType | None = ..., compilation_event: tensorflow.compiler.xla.xla_data_pb2.CompilationEvent.ValueType | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["compilation_event", b"compilation_event", "profile_source", b"profile_source", "profile_type", b"profile_type", "relative_speedup", b"relative_speedup"]) -> None: ... NAME_FIELD_NUMBER: builtins.int ENTRY_COMPUTATION_NAME_FIELD_NUMBER: builtins.int ENTRY_COMPUTATION_ID_FIELD_NUMBER: builtins.int COMPUTATIONS_FIELD_NUMBER: builtins.int HOST_PROGRAM_SHAPE_FIELD_NUMBER: builtins.int ID_FIELD_NUMBER: builtins.int SCHEDULE_FIELD_NUMBER: builtins.int INPUT_OUTPUT_ALIAS_FIELD_NUMBER: builtins.int DYNAMIC_PARAMETER_BINDING_FIELD_NUMBER: builtins.int CROSS_PROGRAM_PREFETCHES_FIELD_NUMBER: builtins.int IS_DYNAMIC_FIELD_NUMBER: builtins.int SPMD_OUTPUT_SHARDING_FIELD_NUMBER: builtins.int SPMD_PARAMETERS_SHARDINGS_FIELD_NUMBER: builtins.int USE_AUTO_SPMD_PARTITIONING_FIELD_NUMBER: builtins.int PROFILE_INFO_FIELD_NUMBER: builtins.int DEVICE_ASSIGNMENT_FIELD_NUMBER: builtins.int name: builtins.str entry_computation_name: builtins.str entry_computation_id: builtins.int @property def computations(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___HloComputationProto]: """The array of computations is always in a valid dependency order, where callees appear before their callers. """ @property def host_program_shape(self) -> tensorflow.compiler.xla.xla_data_pb2.ProgramShapeProto: """The host program shape (with layout) of the entry computation.""" id: builtins.int """The id of this module.""" @property def schedule(self) -> global___HloScheduleProto: """The schedule for this module.""" @property def input_output_alias(self) -> global___HloInputOutputAliasProto: """Describes alias information between inputs and outputs.""" @property def dynamic_parameter_binding(self) -> global___DynamicParameterBindingProto: ... @property def cross_program_prefetches(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___CrossProgramPrefetch]: ... is_dynamic: builtins.bool """True if the module contains dynamic computation.""" @property def spmd_output_sharding(self) -> tensorflow.compiler.xla.xla_data_pb2.OpSharding: ... @property def spmd_parameters_shardings(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[tensorflow.compiler.xla.xla_data_pb2.OpSharding]: ... use_auto_spmd_partitioning: builtins.bool """Uses AutoSharding pass or not.""" @property def profile_info(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___HloModuleProto.ProfileInfo]: """Profile information for the HLO module.""" @property def device_assignment(self) -> tensorflow.compiler.xla.xla_data_pb2.DeviceAssignmentProto: """DeviceAssignment object information.""" def __init__( self, *, name: builtins.str | None = ..., entry_computation_name: builtins.str | None = ..., entry_computation_id: builtins.int | None = ..., computations: collections.abc.Iterable[global___HloComputationProto] | None = ..., host_program_shape: tensorflow.compiler.xla.xla_data_pb2.ProgramShapeProto | None = ..., id: builtins.int | None = ..., schedule: global___HloScheduleProto | None = ..., input_output_alias: global___HloInputOutputAliasProto | None = ..., dynamic_parameter_binding: global___DynamicParameterBindingProto | None = ..., cross_program_prefetches: collections.abc.Iterable[global___CrossProgramPrefetch] | None = ..., is_dynamic: builtins.bool | None = ..., spmd_output_sharding: tensorflow.compiler.xla.xla_data_pb2.OpSharding | None = ..., spmd_parameters_shardings: collections.abc.Iterable[tensorflow.compiler.xla.xla_data_pb2.OpSharding] | None = ..., use_auto_spmd_partitioning: builtins.bool | None = ..., profile_info: collections.abc.Iterable[global___HloModuleProto.ProfileInfo] | None = ..., device_assignment: tensorflow.compiler.xla.xla_data_pb2.DeviceAssignmentProto | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["device_assignment", b"device_assignment", "dynamic_parameter_binding", b"dynamic_parameter_binding", "host_program_shape", b"host_program_shape", "input_output_alias", b"input_output_alias", "schedule", b"schedule", "spmd_output_sharding", b"spmd_output_sharding"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["computations", b"computations", "cross_program_prefetches", b"cross_program_prefetches", "device_assignment", b"device_assignment", "dynamic_parameter_binding", b"dynamic_parameter_binding", "entry_computation_id", b"entry_computation_id", "entry_computation_name", b"entry_computation_name", "host_program_shape", b"host_program_shape", "id", b"id", "input_output_alias", b"input_output_alias", "is_dynamic", b"is_dynamic", "name", b"name", "profile_info", b"profile_info", "schedule", b"schedule", "spmd_output_sharding", b"spmd_output_sharding", "spmd_parameters_shardings", b"spmd_parameters_shardings", "use_auto_spmd_partitioning", b"use_auto_spmd_partitioning"]) -> None: ... global___HloModuleProto = HloModuleProto @typing_extensions.final class LogicalBufferProto(google.protobuf.message.Message): """Serialization of LogicalBuffer.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class Location(google.protobuf.message.Message): """Location represents an instruction and its shape index, which uniquely identifies a point where a buffer is needed. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor INSTRUCTION_NAME_FIELD_NUMBER: builtins.int INSTRUCTION_ID_FIELD_NUMBER: builtins.int SHAPE_INDEX_FIELD_NUMBER: builtins.int instruction_name: builtins.str """TODO(b/239098765): Remove instruction_name and computation_name.""" instruction_id: builtins.int @property def shape_index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ... def __init__( self, *, instruction_name: builtins.str | None = ..., instruction_id: builtins.int | None = ..., shape_index: collections.abc.Iterable[builtins.int] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["instruction_id", b"instruction_id", "instruction_name", b"instruction_name", "shape_index", b"shape_index"]) -> None: ... ID_FIELD_NUMBER: builtins.int SIZE_FIELD_NUMBER: builtins.int DEFINED_AT_FIELD_NUMBER: builtins.int COLOR_FIELD_NUMBER: builtins.int id: builtins.int size: builtins.int @property def defined_at(self) -> global___LogicalBufferProto.Location: """The location where the buffer is defined.""" color: builtins.int def __init__( self, *, id: builtins.int | None = ..., size: builtins.int | None = ..., defined_at: global___LogicalBufferProto.Location | None = ..., color: builtins.int | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["defined_at", b"defined_at"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["color", b"color", "defined_at", b"defined_at", "id", b"id", "size", b"size"]) -> None: ... global___LogicalBufferProto = LogicalBufferProto @typing_extensions.final class BufferAllocationProto(google.protobuf.message.Message): """Serialization of BufferAllocation.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class Assigned(google.protobuf.message.Message): """Assigned represents a single LogicalBuffer that is assigned to this BufferAllocation. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor LOGICAL_BUFFER_ID_FIELD_NUMBER: builtins.int OFFSET_FIELD_NUMBER: builtins.int SIZE_FIELD_NUMBER: builtins.int logical_buffer_id: builtins.int offset: builtins.int size: builtins.int def __init__( self, *, logical_buffer_id: builtins.int | None = ..., offset: builtins.int | None = ..., size: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["logical_buffer_id", b"logical_buffer_id", "offset", b"offset", "size", b"size"]) -> None: ... INDEX_FIELD_NUMBER: builtins.int SIZE_FIELD_NUMBER: builtins.int IS_THREAD_LOCAL_FIELD_NUMBER: builtins.int IS_TUPLE_FIELD_NUMBER: builtins.int IS_ENTRY_COMPUTATION_PARAMETER_FIELD_NUMBER: builtins.int IS_CONSTANT_FIELD_NUMBER: builtins.int PARAMETER_NUMBER_FIELD_NUMBER: builtins.int PARAMETER_SHAPE_INDEX_FIELD_NUMBER: builtins.int MAYBE_LIVE_OUT_FIELD_NUMBER: builtins.int COLOR_FIELD_NUMBER: builtins.int ASSIGNED_FIELD_NUMBER: builtins.int index: builtins.int size: builtins.int is_thread_local: builtins.bool is_tuple: builtins.bool is_entry_computation_parameter: builtins.bool is_constant: builtins.bool parameter_number: builtins.int @property def parameter_shape_index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ... maybe_live_out: builtins.bool color: builtins.int @property def assigned(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___BufferAllocationProto.Assigned]: ... def __init__( self, *, index: builtins.int | None = ..., size: builtins.int | None = ..., is_thread_local: builtins.bool | None = ..., is_tuple: builtins.bool | None = ..., is_entry_computation_parameter: builtins.bool | None = ..., is_constant: builtins.bool | None = ..., parameter_number: builtins.int | None = ..., parameter_shape_index: collections.abc.Iterable[builtins.int] | None = ..., maybe_live_out: builtins.bool | None = ..., color: builtins.int | None = ..., assigned: collections.abc.Iterable[global___BufferAllocationProto.Assigned] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["assigned", b"assigned", "color", b"color", "index", b"index", "is_constant", b"is_constant", "is_entry_computation_parameter", b"is_entry_computation_parameter", "is_thread_local", b"is_thread_local", "is_tuple", b"is_tuple", "maybe_live_out", b"maybe_live_out", "parameter_number", b"parameter_number", "parameter_shape_index", b"parameter_shape_index", "size", b"size"]) -> None: ... global___BufferAllocationProto = BufferAllocationProto @typing_extensions.final class HeapSimulatorTrace(google.protobuf.message.Message): """A trace of a HeapSimulator run.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class Event(google.protobuf.message.Message): """The trace includes a list of events, where each event describes one action performed by the heap simulator. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor class _Kind: ValueType = typing.NewType("ValueType", builtins.int) V: typing_extensions.TypeAlias = ValueType class _KindEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[HeapSimulatorTrace.Event._Kind.ValueType], builtins.type): DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor ALLOC: HeapSimulatorTrace.Event._Kind.ValueType # 0 """A memory region was allocated for the buffer.""" FREE: HeapSimulatorTrace.Event._Kind.ValueType # 1 """A memory region was freed for the buffer.""" SHARE_WITH: HeapSimulatorTrace.Event._Kind.ValueType # 2 """A buffer was shared with another (canonical) buffer. This is similar to ALLOC, except that instead of allocating a new region of memory, the memory region of the canonical buffer is directly re-used. Multiple buffers may share with the same canonical buffer. The lifetime of the canonical buffer is extended to the union of all lifetimes. """ class Kind(_Kind, metaclass=_KindEnumTypeWrapper): ... ALLOC: HeapSimulatorTrace.Event.Kind.ValueType # 0 """A memory region was allocated for the buffer.""" FREE: HeapSimulatorTrace.Event.Kind.ValueType # 1 """A memory region was freed for the buffer.""" SHARE_WITH: HeapSimulatorTrace.Event.Kind.ValueType # 2 """A buffer was shared with another (canonical) buffer. This is similar to ALLOC, except that instead of allocating a new region of memory, the memory region of the canonical buffer is directly re-used. Multiple buffers may share with the same canonical buffer. The lifetime of the canonical buffer is extended to the union of all lifetimes. """ KIND_FIELD_NUMBER: builtins.int BUFFER_ID_FIELD_NUMBER: builtins.int COMPUTATION_NAME_FIELD_NUMBER: builtins.int INSTRUCTION_NAME_FIELD_NUMBER: builtins.int SHARE_WITH_CANONICAL_ID_FIELD_NUMBER: builtins.int kind: global___HeapSimulatorTrace.Event.Kind.ValueType buffer_id: builtins.int """The id of the LogicalBuffer that the event applies to.""" computation_name: builtins.str """The HloInstruction that the simulation was processing that caused this event to occur, identified by its computation and instruction name. E.g. buffers defined by instruction A are allocated when processing A. """ instruction_name: builtins.str share_with_canonical_id: builtins.int """The id of the canonical LogicalBuffer that the buffer shares with. Only set for SHARE_WITH events. """ def __init__( self, *, kind: global___HeapSimulatorTrace.Event.Kind.ValueType | None = ..., buffer_id: builtins.int | None = ..., computation_name: builtins.str | None = ..., instruction_name: builtins.str | None = ..., share_with_canonical_id: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["buffer_id", b"buffer_id", "computation_name", b"computation_name", "instruction_name", b"instruction_name", "kind", b"kind", "share_with_canonical_id", b"share_with_canonical_id"]) -> None: ... EVENTS_FIELD_NUMBER: builtins.int WHOLE_MODULE_SIMULATION_FIELD_NUMBER: builtins.int BUFFER_ALLOCATION_INDEX_FIELD_NUMBER: builtins.int @property def events(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___HeapSimulatorTrace.Event]: ... whole_module_simulation: builtins.bool buffer_allocation_index: builtins.int def __init__( self, *, events: collections.abc.Iterable[global___HeapSimulatorTrace.Event] | None = ..., whole_module_simulation: builtins.bool | None = ..., buffer_allocation_index: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["buffer_allocation_index", b"buffer_allocation_index", "events", b"events", "whole_module_simulation", b"whole_module_simulation"]) -> None: ... global___HeapSimulatorTrace = HeapSimulatorTrace @typing_extensions.final class HloModuleGroupProto(google.protobuf.message.Message): """An abstraction representing a set of HLO module built to run concurrently across different devices. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor NAME_FIELD_NUMBER: builtins.int HLO_MODULES_FIELD_NUMBER: builtins.int name: builtins.str @property def hlo_modules(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___HloModuleProto]: ... def __init__( self, *, name: builtins.str | None = ..., hlo_modules: collections.abc.Iterable[global___HloModuleProto] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["hlo_modules", b"hlo_modules", "name", b"name"]) -> None: ... global___HloModuleGroupProto = HloModuleGroupProto @typing_extensions.final class BufferAssignmentProto(google.protobuf.message.Message): """Serialization of BufferAssignment.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class BufferAlias(google.protobuf.message.Message): """Alias represents a source LogicalBuffer, and the buffer location that aliases it. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor SOURCE_BUFFER_ID_FIELD_NUMBER: builtins.int LOCATION_FIELD_NUMBER: builtins.int source_buffer_id: builtins.int @property def location(self) -> global___LogicalBufferProto.Location: ... def __init__( self, *, source_buffer_id: builtins.int | None = ..., location: global___LogicalBufferProto.Location | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["location", b"location"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["location", b"location", "source_buffer_id", b"source_buffer_id"]) -> None: ... LOGICAL_BUFFERS_FIELD_NUMBER: builtins.int BUFFER_ALIASES_FIELD_NUMBER: builtins.int BUFFER_ALLOCATIONS_FIELD_NUMBER: builtins.int HEAP_SIMULATOR_TRACES_FIELD_NUMBER: builtins.int @property def logical_buffers(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___LogicalBufferProto]: ... @property def buffer_aliases(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___BufferAssignmentProto.BufferAlias]: ... @property def buffer_allocations(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___BufferAllocationProto]: ... @property def heap_simulator_traces(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___HeapSimulatorTrace]: ... def __init__( self, *, logical_buffers: collections.abc.Iterable[global___LogicalBufferProto] | None = ..., buffer_aliases: collections.abc.Iterable[global___BufferAssignmentProto.BufferAlias] | None = ..., buffer_allocations: collections.abc.Iterable[global___BufferAllocationProto] | None = ..., heap_simulator_traces: collections.abc.Iterable[global___HeapSimulatorTrace] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["buffer_aliases", b"buffer_aliases", "buffer_allocations", b"buffer_allocations", "heap_simulator_traces", b"heap_simulator_traces", "logical_buffers", b"logical_buffers"]) -> None: ... global___BufferAssignmentProto = BufferAssignmentProto @typing_extensions.final class HloProto(google.protobuf.message.Message): """Grouping message that contains all of the information above.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor HLO_MODULE_FIELD_NUMBER: builtins.int BUFFER_ASSIGNMENT_FIELD_NUMBER: builtins.int @property def hlo_module(self) -> global___HloModuleProto: ... @property def buffer_assignment(self) -> global___BufferAssignmentProto: ... def __init__( self, *, hlo_module: global___HloModuleProto | None = ..., buffer_assignment: global___BufferAssignmentProto | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["buffer_assignment", b"buffer_assignment", "hlo_module", b"hlo_module"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["buffer_assignment", b"buffer_assignment", "hlo_module", b"hlo_module"]) -> None: ... global___HloProto = HloProto @typing_extensions.final class HloSnapshot(google.protobuf.message.Message): """Encapsulates HloProto together with the arguments, result, and execution_platform. This message is used for purposes such as analysis/replay/file-storage. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor HLO_FIELD_NUMBER: builtins.int ARGUMENTS_FIELD_NUMBER: builtins.int RESULT_FIELD_NUMBER: builtins.int EXECUTION_PLATFORM_FIELD_NUMBER: builtins.int @property def hlo(self) -> global___HloProto: """The hlo graph.""" @property def arguments(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[tensorflow.compiler.xla.xla_data_pb2.LiteralProto]: """The arguments passed to the graph.""" @property def result(self) -> tensorflow.compiler.xla.xla_data_pb2.LiteralProto: """The result of the graph.""" execution_platform: builtins.str """The name of the platform used to run the graph.""" def __init__( self, *, hlo: global___HloProto | None = ..., arguments: collections.abc.Iterable[tensorflow.compiler.xla.xla_data_pb2.LiteralProto] | None = ..., result: tensorflow.compiler.xla.xla_data_pb2.LiteralProto | None = ..., execution_platform: builtins.str | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["hlo", b"hlo", "result", b"result"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["arguments", b"arguments", "execution_platform", b"execution_platform", "hlo", b"hlo", "result", b"result"]) -> None: ... global___HloSnapshot = HloSnapshot @typing_extensions.final class HloModuleMetadataProto(google.protobuf.message.Message): """Metadata for an HLO module. Dumped after HLO passes and before LLO lowering with filename module_####.metadata.textproto, where #### is canonical_module_id. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor CANONICAL_MODULE_ID_FIELD_NUMBER: builtins.int MODULE_GROUP_NAME_FIELD_NUMBER: builtins.int ORIGINAL_MODULE_ID_FIELD_NUMBER: builtins.int PARTITIONED_MODULE_IDS_FIELD_NUMBER: builtins.int PASS_METADATA_FIELD_NUMBER: builtins.int canonical_module_id: builtins.int """Uniquely identifies an HloModuleMetadata. Equal to the first unique_id of the module (a module may go through multiple unique_ids). If a module is partitioned into multiple modules, those modules will each have a new HloModuleMetadata with a different canonical_module_id. """ module_group_name: builtins.str """Name of the module group that the module is part of.""" original_module_id: builtins.int """The canonical module id of the module that this one is partitioned from, if applicable. """ @property def partitioned_module_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: """The canonical module ids of the modules that this one is partitioned into, if applicable. """ @property def pass_metadata(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___HloPassMetadata]: """Metadata for the HLO passes that are run on the module.""" def __init__( self, *, canonical_module_id: builtins.int | None = ..., module_group_name: builtins.str | None = ..., original_module_id: builtins.int | None = ..., partitioned_module_ids: collections.abc.Iterable[builtins.int] | None = ..., pass_metadata: collections.abc.Iterable[global___HloPassMetadata] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["canonical_module_id", b"canonical_module_id", "module_group_name", b"module_group_name", "original_module_id", b"original_module_id", "partitioned_module_ids", b"partitioned_module_ids", "pass_metadata", b"pass_metadata"]) -> None: ... global___HloModuleMetadataProto = HloModuleMetadataProto @typing_extensions.final class HloPassMetadata(google.protobuf.message.Message): """Metadata for one run of an HLO pass on a module. Provides more information when processing debug dumps of HloProtos about the order of HLO passes and various other stats like duration. `pass_id` may also be used to identify a particular run of a pass in debug info that propagates through stages of compilation. """ DESCRIPTOR: google.protobuf.descriptor.Descriptor PASS_ID_FIELD_NUMBER: builtins.int PASS_NAME_FIELD_NUMBER: builtins.int PIPELINE_NAME_FIELD_NUMBER: builtins.int DUMP_FILENAMES_FIELD_NUMBER: builtins.int MODULE_CHANGED_FIELD_NUMBER: builtins.int MODULE_ID_FIELD_NUMBER: builtins.int MODULE_GROUP_MODULE_IDS_FIELD_NUMBER: builtins.int START_TIMESTAMP_USEC_FIELD_NUMBER: builtins.int END_TIMESTAMP_USEC_FIELD_NUMBER: builtins.int pass_id: builtins.int """For a given module, pass_id uniquely identifies a run of an HLO pass on that module. Note that a pass_id may not always refer to the same pass because the order of passes during compilation may change. For finding metadata for a particular pass, pass_name and pipeline_name would be more reliable, although note that they may not be unique. """ pass_name: builtins.str pipeline_name: builtins.str @property def dump_filenames(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """Filenames of the dumps of the module after this pass ran. Module may be dumped in multiple formats, and the order of formats in this field will stay consistent across passes. """ module_changed: builtins.bool """Return value of pass.Run(). True if this pass changed the module, or, in the case where the module was run through this pass as part of a module group, true if this pass changed any module in the same module group. """ module_id: builtins.int """The unique_id of the module that this pass is run on. May be different from the canonical_module_id of the HloModuleMetadata that this HloPassMetadata is inside. """ @property def module_group_module_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: """If the module went through this pass as part of a module group, this is set as the ids of all the modules in the module group. Empty otherwise. """ start_timestamp_usec: builtins.int """Timestamp before and after the pass is run. Note they may be equal.""" end_timestamp_usec: builtins.int def __init__( self, *, pass_id: builtins.int | None = ..., pass_name: builtins.str | None = ..., pipeline_name: builtins.str | None = ..., dump_filenames: collections.abc.Iterable[builtins.str] | None = ..., module_changed: builtins.bool | None = ..., module_id: builtins.int | None = ..., module_group_module_ids: collections.abc.Iterable[builtins.int] | None = ..., start_timestamp_usec: builtins.int | None = ..., end_timestamp_usec: builtins.int | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["dump_filenames", b"dump_filenames", "end_timestamp_usec", b"end_timestamp_usec", "module_changed", b"module_changed", "module_group_module_ids", b"module_group_module_ids", "module_id", b"module_id", "pass_id", b"pass_id", "pass_name", b"pass_name", "pipeline_name", b"pipeline_name", "start_timestamp_usec", b"start_timestamp_usec"]) -> None: ... global___HloPassMetadata = HloPassMetadata @typing_extensions.final class EntryFunctionAttributes(google.protobuf.message.Message): """Encodes attributes for an entry function.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor @typing_extensions.final class ShapeIndex(google.protobuf.message.Message): """Acts as the underlying container for an xla::ShapeIndex.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor INDICES_FIELD_NUMBER: builtins.int @property def indices(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ... def __init__( self, *, indices: collections.abc.Iterable[builtins.int] | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["indices", b"indices"]) -> None: ... @typing_extensions.final class BufferParameterAttributes(google.protobuf.message.Message): """Encodes attributes for a single buffer parameter.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor LMHLO_PARAMS_FIELD_NUMBER: builtins.int LMHLO_PARAMS_PRESENT_FIELD_NUMBER: builtins.int LMHLO_PARAM_SHAPE_INDEX_FIELD_NUMBER: builtins.int LMHLO_CONSTANT_NAME_FIELD_NUMBER: builtins.int LMHLO_MUST_ALIAS_FIELD_NUMBER: builtins.int LMHLO_OUTPUT_INDEX_FIELD_NUMBER: builtins.int lmhlo_params: builtins.int """Represents an lmhlo.params function argument attribute.""" lmhlo_params_present: builtins.bool """TODO(hanbinyoon): Deprecate when optional fields are available in proto3 (Protocol Buffers v3.15.0). """ @property def lmhlo_param_shape_index(self) -> global___EntryFunctionAttributes.ShapeIndex: """Represents an lmhlo.param_shape_index function argument attribute.""" lmhlo_constant_name: builtins.str """Represents an lmhlo.constant_name function argument attribute.""" lmhlo_must_alias: builtins.bool """Represents an lmhlo.must_alias function argument attribute.""" @property def lmhlo_output_index(self) -> global___EntryFunctionAttributes.ShapeIndex: """Represents an lmhlo.params function argument attribute.""" def __init__( self, *, lmhlo_params: builtins.int | None = ..., lmhlo_params_present: builtins.bool | None = ..., lmhlo_param_shape_index: global___EntryFunctionAttributes.ShapeIndex | None = ..., lmhlo_constant_name: builtins.str | None = ..., lmhlo_must_alias: builtins.bool | None = ..., lmhlo_output_index: global___EntryFunctionAttributes.ShapeIndex | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["lmhlo_output_index", b"lmhlo_output_index", "lmhlo_param_shape_index", b"lmhlo_param_shape_index"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["lmhlo_constant_name", b"lmhlo_constant_name", "lmhlo_must_alias", b"lmhlo_must_alias", "lmhlo_output_index", b"lmhlo_output_index", "lmhlo_param_shape_index", b"lmhlo_param_shape_index", "lmhlo_params", b"lmhlo_params", "lmhlo_params_present", b"lmhlo_params_present"]) -> None: ... BUFFERS_FIELD_NUMBER: builtins.int RESULT_XLA_SHAPE_FIELD_NUMBER: builtins.int @property def buffers(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___EntryFunctionAttributes.BufferParameterAttributes]: ... result_xla_shape: builtins.str """xla::Shape in string format.""" def __init__( self, *, buffers: collections.abc.Iterable[global___EntryFunctionAttributes.BufferParameterAttributes] | None = ..., result_xla_shape: builtins.str | None = ..., ) -> None: ... def ClearField(self, field_name: typing_extensions.Literal["buffers", b"buffers", "result_xla_shape", b"result_xla_shape"]) -> None: ... global___EntryFunctionAttributes = EntryFunctionAttributes @typing_extensions.final class XlaRuntimeExecutableProto(google.protobuf.message.Message): """Encodes the underlying Xla runtime executable compiled from the XLA module.""" DESCRIPTOR: google.protobuf.descriptor.Descriptor HLO_MODULE_PROTO_FIELD_NUMBER: builtins.int OBJ_FILE_FIELD_NUMBER: builtins.int MLIR_MODULE_FIELD_NUMBER: builtins.int @property def hlo_module_proto(self) -> global___HloModuleProto: ... obj_file: builtins.bytes """TODO(b/232263665)): Serialized executable has to know what APIs it has to be linked with, including the version. For example Gpu executable must be linked with a runtime layer that abstracts over CUDA. Serialized object file compiled from the XLA module. """ mlir_module: builtins.str """Serialized MLIR module corresponding to compiled object file.""" def __init__( self, *, hlo_module_proto: global___HloModuleProto | None = ..., obj_file: builtins.bytes | None = ..., mlir_module: builtins.str | None = ..., ) -> None: ... def HasField(self, field_name: typing_extensions.Literal["hlo_module_proto", b"hlo_module_proto"]) -> builtins.bool: ... def ClearField(self, field_name: typing_extensions.Literal["hlo_module_proto", b"hlo_module_proto", "mlir_module", b"mlir_module", "obj_file", b"obj_file"]) -> None: ... global___XlaRuntimeExecutableProto = XlaRuntimeExecutableProto