basedpyright 1.36.2 → 1.37.0-961d59bacff1958014fe4e9ed2ef3116fe2266e2

package/dist/typeshed-fallback/stubs/tensorflow/tensorflow/compiler/xla/service/hlo_pb2.pyi

@@ -0,0 +1,2113 @@
+"""
+@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.any_pb2
+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<StridedMemrefView>()
+        .Attr<float>("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<StridedMemrefView>()
+    .Attr<float>("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.final
+class HloInstructionProto(google.protobuf.message.Message):
+    """Serialization of HloInstruction.
+    Next ID: 90
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.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.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_EXECUTION_THREAD_FIELD_NUMBER: builtins.int
+    K_FIELD_NUMBER: builtins.int
+    LARGEST_FIELD_NUMBER: builtins.int
+    STATISTICS_VIZ_FIELD_NUMBER: builtins.int
+    DOT_SPARSITY_FIELD_NUMBER: builtins.int
+    COLLECTIVE_DEVICE_LIST_FIELD_NUMBER: builtins.int
+    ORIGINAL_VALUE_FIELD_NUMBER: builtins.int
+    IS_COMPOSITE_FIELD_NUMBER: builtins.int
+    name: builtins.str
+    opcode: builtins.str
+    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."""
+    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
+    exponent_bits: builtins.int
+    """The bit sizes for a reduce-precision operation."""
+    mantissa_bits: builtins.int
+    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.
+    """
+    fft_type: tensorflow.compiler.xla.xla_data_pb2.FftType.ValueType
+    """FFT type (FFT, IFFT, etc)."""
+    comparison_direction: builtins.str
+    """Comparison direction only used for kCompare."""
+    id: builtins.int
+    """The id of this instruction."""
+    backend_config: builtins.bytes
+    """Backend configuration for the instruction. Has backend-specific meaning."""
+    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."""
+    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.
+    """
+    custom_call_has_side_effect: builtins.bool
+    """Whether the kCustomCall instruction has side-effects, only present for
+    kCustomCall.
+    """
+    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.
+    """
+    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_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.
+    """
+    k: builtins.int
+    """Represents the K value for top-k."""
+    largest: builtins.bool
+    """Represents the largest flag for top-k."""
+    is_composite: builtins.bool
+    """Specifies if a call instruction is a composite."""
+    @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."""
+
+    @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."""
+
+    @property
+    def slice_dimensions(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___HloInstructionProto.SliceDimensions]: ...
+    @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.
+        """
+
+    @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"""
+
+    @property
+    def fft_length(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """FFT length."""
+
+    @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]: ...
+    @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: ...
+    @property
+    def replica_groups(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[tensorflow.compiler.xla.xla_data_pb2.ReplicaGroup]:
+        """Deprecated, but keeping for backward compatibility.
+        Use collective_device_list. Cross replica op fields.
+        """
+
+    @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: ...
+    @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."""
+
+    @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.
+        """
+
+    @property
+    def frontend_attributes(self) -> tensorflow.compiler.xla.xla_data_pb2.FrontendAttributes:
+        """Frontend attributes to pass to the XLA backend."""
+
+    @property
+    def statistics_viz(self) -> tensorflow.compiler.xla.xla_data_pb2.StatisticsViz:
+        """Represents the information for tracking propagation of values within HLO
+        graph.
+        """
+
+    @property
+    def dot_sparsity(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[
+        tensorflow.compiler.xla.xla_data_pb2.SparsityDescriptor
+    ]:
+        """Sparsity descriptor for dot operation."""
+
+    @property
+    def collective_device_list(self) -> tensorflow.compiler.xla.xla_data_pb2.CollectiveDeviceListProto:
+        """Represents the list of devices that participate in a collective operation."""
+
+    @property
+    def original_value(self) -> tensorflow.compiler.xla.xla_data_pb2.OriginalValueProto:
+        """For HLO value tracking."""
+
+    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_execution_thread: builtins.str | None = ...,
+        k: builtins.int | None = ...,
+        largest: builtins.bool | None = ...,
+        statistics_viz: tensorflow.compiler.xla.xla_data_pb2.StatisticsViz | None = ...,
+        dot_sparsity: collections.abc.Iterable[tensorflow.compiler.xla.xla_data_pb2.SparsityDescriptor] | None = ...,
+        collective_device_list: tensorflow.compiler.xla.xla_data_pb2.CollectiveDeviceListProto | None = ...,
+        original_value: tensorflow.compiler.xla.xla_data_pb2.OriginalValueProto | None = ...,
+        is_composite: builtins.bool | None = ...,
+    ) -> None: ...
+    def HasField(
+        self,
+        field_name: typing.Literal[
+            "cholesky_options",
+            b"cholesky_options",
+            "collective_device_list",
+            b"collective_device_list",
+            "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",
+            "original_value",
+            b"original_value",
+            "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",
+            "statistics_viz",
+            b"statistics_viz",
+            "triangular_solve_options",
+            b"triangular_solve_options",
+            "window",
+            b"window",
+        ],
+    ) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "all_reduce_id",
+            b"all_reduce_id",
+            "async_execution_thread",
+            b"async_execution_thread",
+            "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",
+            "collective_device_list",
+            b"collective_device_list",
+            "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",
+            "dot_sparsity",
+            b"dot_sparsity",
+            "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_composite",
+            b"is_composite",
+            "is_cross_program_prefetch",
+            b"is_cross_program_prefetch",
+            "is_host_transfer",
+            b"is_host_transfer",
+            "is_stable",
+            b"is_stable",
+            "k",
+            b"k",
+            "largest",
+            b"largest",
+            "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",
+            "original_value",
+            b"original_value",
+            "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",
+            "statistics_viz",
+            b"statistics_viz",
+            "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.Literal["optional_cross_program_prefetch_index", b"optional_cross_program_prefetch_index"]
+    ) -> typing.Literal["cross_program_prefetch_index"] | None: ...
+
+global___HloInstructionProto = HloInstructionProto
+
+@typing.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
+    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."""
+    @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."""
+
+    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.Literal["program_shape", b"program_shape"]) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing.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.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.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.Literal["instruction_ids", b"instruction_ids"]) -> None: ...
+
+    @typing.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.Literal["value", b"value"]) -> builtins.bool: ...
+        def ClearField(self, field_name: typing.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.Literal["sequences", b"sequences"]) -> None: ...
+
+global___HloScheduleProto = HloScheduleProto
+
+@typing.final
+class HloInputOutputAliasProto(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.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 parameter'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
+        parameter_number: builtins.int
+        """Number of the parameter in entry computation."""
+        kind: global___Kind.ValueType
+        """The kind of alias to be setup."""
+        @property
+        def output_shape_index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+            """ShapeIndex of the root hlo."""
+
+        @property
+        def parameter_shape_index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+            """ShapeIndex of the parameter instruction."""
+
+        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.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.Literal["entries", b"entries"]) -> None: ...
+
+global___HloInputOutputAliasProto = HloInputOutputAliasProto
+
+@typing.final
+class HloBufferDonorProto(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.final
+    class BufferDonorEntryProto(google.protobuf.message.Message):
+        """The following proto describes an input (described by parameter number and a
+        ShapeIndex of the parameter) that can donate its butter to any output
+        tensor. It is similar to HloInputOutputAliasProto, but without a paired
+        output. For example:
+
+        entry = {
+         parameter_number=0,
+         parameter_shape_index={1, 2},
+        }
+
+        This entry indicates that the first parameter's {1, 2} element can donate
+        its buffer.
+        """
+
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        PARAMETER_NUMBER_FIELD_NUMBER: builtins.int
+        PARAMETER_SHAPE_INDEX_FIELD_NUMBER: builtins.int
+        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."""
+
+        def __init__(
+            self,
+            *,
+            parameter_number: builtins.int | None = ...,
+            parameter_shape_index: collections.abc.Iterable[builtins.int] | None = ...,
+        ) -> None: ...
+        def ClearField(
+            self,
+            field_name: typing.Literal[
+                "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___HloBufferDonorProto.BufferDonorEntryProto
+    ]: ...
+    def __init__(
+        self, *, entries: collections.abc.Iterable[global___HloBufferDonorProto.BufferDonorEntryProto] | None = ...
+    ) -> None: ...
+    def ClearField(self, field_name: typing.Literal["entries", b"entries"]) -> None: ...
+
+global___HloBufferDonorProto = HloBufferDonorProto
+
+@typing.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
+    offset: builtins.int
+    @property
+    def index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[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.Literal["index", b"index", "offset", b"offset", "parameter", b"parameter"]
+    ) -> None: ...
+
+global___CrossProgramPrefetch = CrossProgramPrefetch
+
+@typing.final
+class StackFrameIndexProto(google.protobuf.message.Message):
+    """Serialization of stack frames index representations.
+    Stack frames index presented in four flat arrays:
+    1. File names array.
+    2. Function names array.
+    3. File location array.
+    4. Frame array.
+    All reference ids in sub-protos are 1-based positions of the
+    entity in the flat array.
+    Ids are 1-based to keep 0 value as representation of non-set property.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.final
+    class FileLocation(google.protobuf.message.Message):
+        """Serialization of file position."""
+
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        FILE_NAME_ID_FIELD_NUMBER: builtins.int
+        FUNCTION_NAME_ID_FIELD_NUMBER: builtins.int
+        LINE_FIELD_NUMBER: builtins.int
+        COLUMN_FIELD_NUMBER: builtins.int
+        file_name_id: builtins.int
+        """1-based position of file name."""
+        function_name_id: builtins.int
+        """1-based position of function name."""
+        line: builtins.int
+        """Line number."""
+        column: builtins.int
+        """Column number."""
+        def __init__(
+            self,
+            *,
+            file_name_id: builtins.int | None = ...,
+            function_name_id: builtins.int | None = ...,
+            line: builtins.int | None = ...,
+            column: builtins.int | None = ...,
+        ) -> None: ...
+        def ClearField(
+            self,
+            field_name: typing.Literal[
+                "column", b"column", "file_name_id", b"file_name_id", "function_name_id", b"function_name_id", "line", b"line"
+            ],
+        ) -> None: ...
+
+    @typing.final
+    class StackFrame(google.protobuf.message.Message):
+        """Serialization of frame."""
+
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        FILE_LOCATION_ID_FIELD_NUMBER: builtins.int
+        PARENT_FRAME_ID_FIELD_NUMBER: builtins.int
+        file_location_id: builtins.int
+        """1-based position of file location."""
+        parent_frame_id: builtins.int
+        """1-based position of the parent frame."""
+        def __init__(
+            self, *, file_location_id: builtins.int | None = ..., parent_frame_id: builtins.int | None = ...
+        ) -> None: ...
+        def ClearField(
+            self, field_name: typing.Literal["file_location_id", b"file_location_id", "parent_frame_id", b"parent_frame_id"]
+        ) -> None: ...
+
+    FILE_NAMES_FIELD_NUMBER: builtins.int
+    FUNCTION_NAMES_FIELD_NUMBER: builtins.int
+    FILE_LOCATIONS_FIELD_NUMBER: builtins.int
+    STACK_FRAMES_FIELD_NUMBER: builtins.int
+    @property
+    def file_names(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]:
+        """Flat index array of file names."""
+
+    @property
+    def function_names(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]:
+        """Flat index array of function names."""
+
+    @property
+    def file_locations(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___StackFrameIndexProto.FileLocation]:
+        """Flat index array of file locations."""
+
+    @property
+    def stack_frames(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___StackFrameIndexProto.StackFrame]:
+        """Flat index array of frames."""
+
+    def __init__(
+        self,
+        *,
+        file_names: collections.abc.Iterable[builtins.str] | None = ...,
+        function_names: collections.abc.Iterable[builtins.str] | None = ...,
+        file_locations: collections.abc.Iterable[global___StackFrameIndexProto.FileLocation] | None = ...,
+        stack_frames: collections.abc.Iterable[global___StackFrameIndexProto.StackFrame] | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "file_locations",
+            b"file_locations",
+            "file_names",
+            b"file_names",
+            "function_names",
+            b"function_names",
+            "stack_frames",
+            b"stack_frames",
+        ],
+    ) -> None: ...
+
+global___StackFrameIndexProto = StackFrameIndexProto
+
+@typing.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
+        FLAGNET: HloModuleProto._ProfileType.ValueType  # 5
+
+    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
+    FLAGNET: HloModuleProto.ProfileType.ValueType  # 5
+
+    @typing.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
+        FINGERPRINT_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."""
+        fingerprint: builtins.str
+        """The fingerprint of the unoptimized module this profile was applied to."""
+        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 = ...,
+            fingerprint: builtins.str | None = ...,
+        ) -> None: ...
+        def ClearField(
+            self,
+            field_name: typing.Literal[
+                "compilation_event",
+                b"compilation_event",
+                "fingerprint",
+                b"fingerprint",
+                "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
+    BUFFER_DONOR_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
+    STACK_FRAME_INDEX_FIELD_NUMBER: builtins.int
+    FRONTEND_ATTRIBUTES_FIELD_NUMBER: builtins.int
+    name: builtins.str
+    entry_computation_name: builtins.str
+    entry_computation_id: builtins.int
+    id: builtins.int
+    """The id of this module."""
+    is_dynamic: builtins.bool
+    """True if the module contains dynamic computation."""
+    use_auto_spmd_partitioning: builtins.bool
+    """Uses AutoSharding pass or not."""
+    @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."""
+
+    @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 buffer_donor(self) -> global___HloBufferDonorProto:
+        """Describes the information of input buffer donors."""
+
+    @property
+    def cross_program_prefetches(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___CrossProgramPrefetch]: ...
+    @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]: ...
+    @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."""
+
+    @property
+    def stack_frame_index(self) -> global___StackFrameIndexProto:
+        """Stack frames index."""
+
+    @property
+    def frontend_attributes(self) -> tensorflow.compiler.xla.xla_data_pb2.FrontendAttributes:
+        """Frontend attributes to pass to the XLA backend."""
+
+    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 = ...,
+        buffer_donor: global___HloBufferDonorProto | 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 = ...,
+        stack_frame_index: global___StackFrameIndexProto | None = ...,
+        frontend_attributes: tensorflow.compiler.xla.xla_data_pb2.FrontendAttributes | None = ...,
+    ) -> None: ...
+    def HasField(
+        self,
+        field_name: typing.Literal[
+            "buffer_donor",
+            b"buffer_donor",
+            "device_assignment",
+            b"device_assignment",
+            "frontend_attributes",
+            b"frontend_attributes",
+            "host_program_shape",
+            b"host_program_shape",
+            "input_output_alias",
+            b"input_output_alias",
+            "schedule",
+            b"schedule",
+            "spmd_output_sharding",
+            b"spmd_output_sharding",
+            "stack_frame_index",
+            b"stack_frame_index",
+        ],
+    ) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "buffer_donor",
+            b"buffer_donor",
+            "computations",
+            b"computations",
+            "cross_program_prefetches",
+            b"cross_program_prefetches",
+            "device_assignment",
+            b"device_assignment",
+            "entry_computation_id",
+            b"entry_computation_id",
+            "entry_computation_name",
+            b"entry_computation_name",
+            "frontend_attributes",
+            b"frontend_attributes",
+            "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",
+            "stack_frame_index",
+            b"stack_frame_index",
+            "use_auto_spmd_partitioning",
+            b"use_auto_spmd_partitioning",
+        ],
+    ) -> None: ...
+
+global___HloModuleProto = HloModuleProto
+
+@typing.final
+class LogicalBufferProto(google.protobuf.message.Message):
+    """Serialization of LogicalBuffer."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.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.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
+    color: builtins.int
+    @property
+    def defined_at(self) -> global___LogicalBufferProto.Location:
+        """The location where the buffer is defined."""
+
+    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.Literal["defined_at", b"defined_at"]) -> builtins.bool: ...
+    def ClearField(
+        self, field_name: typing.Literal["color", b"color", "defined_at", b"defined_at", "id", b"id", "size", b"size"]
+    ) -> None: ...
+
+global___LogicalBufferProto = LogicalBufferProto
+
+@typing.final
+class BufferAllocationProto(google.protobuf.message.Message):
+    """Serialization of BufferAllocation."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.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.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
+    maybe_live_out: builtins.bool
+    color: builtins.int
+    @property
+    def parameter_shape_index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[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.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.final
+class HeapSimulatorTrace(google.protobuf.message.Message):
+    """A trace of a HeapSimulator run."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.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.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
+    whole_module_simulation: builtins.bool
+    buffer_allocation_index: builtins.int
+    @property
+    def events(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___HeapSimulatorTrace.Event]: ...
+    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.Literal[
+            "buffer_allocation_index",
+            b"buffer_allocation_index",
+            "events",
+            b"events",
+            "whole_module_simulation",
+            b"whole_module_simulation",
+        ],
+    ) -> None: ...
+
+global___HeapSimulatorTrace = HeapSimulatorTrace
+
+@typing.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.Literal["hlo_modules", b"hlo_modules", "name", b"name"]) -> None: ...
+
+global___HloModuleGroupProto = HloModuleGroupProto
+
+@typing.final
+class BufferAssignmentProto(google.protobuf.message.Message):
+    """Serialization of BufferAssignment."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.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.Literal["location", b"location"]) -> builtins.bool: ...
+        def ClearField(
+            self, field_name: typing.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.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.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.Literal["buffer_assignment", b"buffer_assignment", "hlo_module", b"hlo_module"]
+    ) -> builtins.bool: ...
+    def ClearField(
+        self, field_name: typing.Literal["buffer_assignment", b"buffer_assignment", "hlo_module", b"hlo_module"]
+    ) -> None: ...
+
+global___HloProto = HloProto
+
+@typing.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
+    execution_platform: builtins.str
+    """The name of the platform used to run the graph."""
+    @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."""
+
+    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.Literal["hlo", b"hlo", "result", b"result"]) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "arguments", b"arguments", "execution_platform", b"execution_platform", "hlo", b"hlo", "result", b"result"
+        ],
+    ) -> None: ...
+
+global___HloSnapshot = HloSnapshot
+
+@typing.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.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.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
+    CUSTOM_METADATA_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
+    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.
+    """
+    start_timestamp_usec: builtins.int
+    """Timestamp before and after the pass is run. Note they may be equal."""
+    end_timestamp_usec: builtins.int
+    @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.
+        """
+
+    @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.
+        """
+
+    @property
+    def custom_metadata(self) -> google.protobuf.any_pb2.Any:
+        """Custom metadata for the pass."""
+
+    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 = ...,
+        custom_metadata: google.protobuf.any_pb2.Any | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["custom_metadata", b"custom_metadata"]) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "custom_metadata",
+            b"custom_metadata",
+            "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