basedpyright 1.36.2 → 1.37.0-961d59bacff1958014fe4e9ed2ef3116fe2266e2

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

@@ -0,0 +1,2681 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+Copyright 2017 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================
+"""
+
+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
+
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+class _PrimitiveType:
+    ValueType = typing.NewType("ValueType", builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+
+class _PrimitiveTypeEnumTypeWrapper(
+    google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_PrimitiveType.ValueType], builtins.type
+):
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    PRIMITIVE_TYPE_INVALID: _PrimitiveType.ValueType  # 0
+    """Invalid primitive type to serve as default."""
+    PRED: _PrimitiveType.ValueType  # 1
+    """Predicates are two-state booleans."""
+    S2: _PrimitiveType.ValueType  # 26
+    """Signed integral values of fixed width."""
+    S4: _PrimitiveType.ValueType  # 21
+    S8: _PrimitiveType.ValueType  # 2
+    S16: _PrimitiveType.ValueType  # 3
+    S32: _PrimitiveType.ValueType  # 4
+    S64: _PrimitiveType.ValueType  # 5
+    U2: _PrimitiveType.ValueType  # 27
+    """Unsigned integral values of fixed width."""
+    U4: _PrimitiveType.ValueType  # 22
+    U8: _PrimitiveType.ValueType  # 6
+    U16: _PrimitiveType.ValueType  # 7
+    U32: _PrimitiveType.ValueType  # 8
+    U64: _PrimitiveType.ValueType  # 9
+    F16: _PrimitiveType.ValueType  # 10
+    """Floating-point values of fixed width.
+
+    Note: if f16s are not natively supported on the device, they will be
+    converted to f16 from f32 at arbirary points in the computation.
+    """
+    F32: _PrimitiveType.ValueType  # 11
+    BF16: _PrimitiveType.ValueType  # 16
+    """Truncated 16 bit floating-point format. This is similar to IEEE's 16 bit
+    floating-point format, but uses 1 bit for the sign, 8 bits for the exponent
+    and 7 bits for the mantissa.
+    """
+    F64: _PrimitiveType.ValueType  # 12
+    F8E5M2: _PrimitiveType.ValueType  # 19
+    """FP8 dtypes, as described in this paper: https://arxiv.org/abs/2209.05433
+
+    F8E5M2 has 5 exponent bits and 2 mantissa bits, and is similar to the
+    existing IEEE types.
+
+    F8E4M3FN has 4 exponent bits and 3 mantissa bits. The "FN" means only
+    Finite and NaN values are supported. Unlike IEEE types, infinities are not
+    supported.  NaN is represented when the exponent and mantissa bits are all
+    1s. All other values are finite.
+
+    F8E4M3B11FNUZ has 4 exponent bits and 3 mantissa bits and a bias of 11. The
+    "FNUZ" means only Finite and NaN values are supported; zero is unsigned.
+    Unlike IEEE types, infinities are not supported.  NaN is represented when
+    the exponent and mantissa bits are all 0s with a sign bit of 1. All other
+    values are finite.
+
+    Support for these dtypes is under development. They do not yet work
+    properly in most cases.
+    TODO(b/259609697): Fully support FP8.
+    """
+    F8E4M3FN: _PrimitiveType.ValueType  # 20
+    F8E4M3B11FNUZ: _PrimitiveType.ValueType  # 23
+    F8E5M2FNUZ: _PrimitiveType.ValueType  # 24
+    """FP8 dtypes, as described in this paper: https://arxiv.org/abs/2206.02915
+
+    F8E5M2FNUZ has 5 exponent bits and 2 mantissa bits.
+    F8E4M3FNUZ has 4 exponent bits and 3 mantissa bits.
+
+    The "FNUZ" means only Finite and NaN values are supported; zero is
+    unsigned. Unlike IEEE types, infinities are not supported.  NaN is
+    represented when the exponent and mantissa bits are all 0s with a sign bit
+    of 1. All other values are finite.
+
+    These differences mean there's an additional exponent value available. To
+    keep the same dynamic range as an IEEE-like FP8 type, the exponent is
+    biased one more than would be expected given the number of exponent bits
+    (8 for Float8E4M3FNUZ and 16 for Float8E5M2FNUZ).
+    """
+    F8E4M3FNUZ: _PrimitiveType.ValueType  # 25
+    C64: _PrimitiveType.ValueType  # 15
+    """Complex values of fixed width.
+    Paired F32 (real, imag), as in std::complex<float>.
+    """
+    C128: _PrimitiveType.ValueType  # 18
+    """Paired F64 (real, imag), as in std::complex<double>."""
+    TUPLE: _PrimitiveType.ValueType  # 13
+    """A tuple is a polymorphic sequence; e.g. a shape that holds different
+    sub-shapes. They are used for things like returning multiple values from a
+    computation; e.g. a computation that returns weights and biases may have a
+    signature that results in a tuple like (f32[784x2000], f32[2000])
+
+    If a shape proto has the tuple element type, it may not have any entries
+    in the dimensions field.
+    """
+    OPAQUE_TYPE: _PrimitiveType.ValueType  # 14
+    """An opaque type used for passing context-specific data to a custom
+    operation. Shapes of this primitive type will have empty dimensions and
+    tuple_shapes fields.
+
+    (OPAQUE would be a better name for this identifier, but that conflicts with
+    a macro defined in windows.h.)
+    """
+    TOKEN: _PrimitiveType.ValueType  # 17
+    """A token type threaded between side-effecting operations. Shapes of this
+    primitive type will have empty dimensions and tuple_shapes fields.
+    """
+
+class PrimitiveType(_PrimitiveType, metaclass=_PrimitiveTypeEnumTypeWrapper):
+    """Primitive types are the individual values that can be held in rectangular
+    multidimensional arrays. A description of the rectangular multidimensional
+    array dimensions / primitive type is given by Shape, below.
+
+    LINT.IfChange
+    """
+
+PRIMITIVE_TYPE_INVALID: PrimitiveType.ValueType  # 0
+"""Invalid primitive type to serve as default."""
+PRED: PrimitiveType.ValueType  # 1
+"""Predicates are two-state booleans."""
+S2: PrimitiveType.ValueType  # 26
+"""Signed integral values of fixed width."""
+S4: PrimitiveType.ValueType  # 21
+S8: PrimitiveType.ValueType  # 2
+S16: PrimitiveType.ValueType  # 3
+S32: PrimitiveType.ValueType  # 4
+S64: PrimitiveType.ValueType  # 5
+U2: PrimitiveType.ValueType  # 27
+"""Unsigned integral values of fixed width."""
+U4: PrimitiveType.ValueType  # 22
+U8: PrimitiveType.ValueType  # 6
+U16: PrimitiveType.ValueType  # 7
+U32: PrimitiveType.ValueType  # 8
+U64: PrimitiveType.ValueType  # 9
+F16: PrimitiveType.ValueType  # 10
+"""Floating-point values of fixed width.
+
+Note: if f16s are not natively supported on the device, they will be
+converted to f16 from f32 at arbirary points in the computation.
+"""
+F32: PrimitiveType.ValueType  # 11
+BF16: PrimitiveType.ValueType  # 16
+"""Truncated 16 bit floating-point format. This is similar to IEEE's 16 bit
+floating-point format, but uses 1 bit for the sign, 8 bits for the exponent
+and 7 bits for the mantissa.
+"""
+F64: PrimitiveType.ValueType  # 12
+F8E5M2: PrimitiveType.ValueType  # 19
+"""FP8 dtypes, as described in this paper: https://arxiv.org/abs/2209.05433
+
+F8E5M2 has 5 exponent bits and 2 mantissa bits, and is similar to the
+existing IEEE types.
+
+F8E4M3FN has 4 exponent bits and 3 mantissa bits. The "FN" means only
+Finite and NaN values are supported. Unlike IEEE types, infinities are not
+supported.  NaN is represented when the exponent and mantissa bits are all
+1s. All other values are finite.
+
+F8E4M3B11FNUZ has 4 exponent bits and 3 mantissa bits and a bias of 11. The
+"FNUZ" means only Finite and NaN values are supported; zero is unsigned.
+Unlike IEEE types, infinities are not supported.  NaN is represented when
+the exponent and mantissa bits are all 0s with a sign bit of 1. All other
+values are finite.
+
+Support for these dtypes is under development. They do not yet work
+properly in most cases.
+TODO(b/259609697): Fully support FP8.
+"""
+F8E4M3FN: PrimitiveType.ValueType  # 20
+F8E4M3B11FNUZ: PrimitiveType.ValueType  # 23
+F8E5M2FNUZ: PrimitiveType.ValueType  # 24
+"""FP8 dtypes, as described in this paper: https://arxiv.org/abs/2206.02915
+
+F8E5M2FNUZ has 5 exponent bits and 2 mantissa bits.
+F8E4M3FNUZ has 4 exponent bits and 3 mantissa bits.
+
+The "FNUZ" means only Finite and NaN values are supported; zero is
+unsigned. Unlike IEEE types, infinities are not supported.  NaN is
+represented when the exponent and mantissa bits are all 0s with a sign bit
+of 1. All other values are finite.
+
+These differences mean there's an additional exponent value available. To
+keep the same dynamic range as an IEEE-like FP8 type, the exponent is
+biased one more than would be expected given the number of exponent bits
+(8 for Float8E4M3FNUZ and 16 for Float8E5M2FNUZ).
+"""
+F8E4M3FNUZ: PrimitiveType.ValueType  # 25
+C64: PrimitiveType.ValueType  # 15
+"""Complex values of fixed width.
+Paired F32 (real, imag), as in std::complex<float>.
+"""
+C128: PrimitiveType.ValueType  # 18
+"""Paired F64 (real, imag), as in std::complex<double>."""
+TUPLE: PrimitiveType.ValueType  # 13
+"""A tuple is a polymorphic sequence; e.g. a shape that holds different
+sub-shapes. They are used for things like returning multiple values from a
+computation; e.g. a computation that returns weights and biases may have a
+signature that results in a tuple like (f32[784x2000], f32[2000])
+
+If a shape proto has the tuple element type, it may not have any entries
+in the dimensions field.
+"""
+OPAQUE_TYPE: PrimitiveType.ValueType  # 14
+"""An opaque type used for passing context-specific data to a custom
+operation. Shapes of this primitive type will have empty dimensions and
+tuple_shapes fields.
+
+(OPAQUE would be a better name for this identifier, but that conflicts with
+a macro defined in windows.h.)
+"""
+TOKEN: PrimitiveType.ValueType  # 17
+"""A token type threaded between side-effecting operations. Shapes of this
+primitive type will have empty dimensions and tuple_shapes fields.
+"""
+global___PrimitiveType = PrimitiveType
+
+class _DimLevelType:
+    ValueType = typing.NewType("ValueType", builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+
+class _DimLevelTypeEnumTypeWrapper(
+    google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_DimLevelType.ValueType], builtins.type
+):
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    DIM_DENSE: _DimLevelType.ValueType  # 0
+    """The corresponding dimension is Dense, every entry is stored."""
+    DIM_COMPRESSED: _DimLevelType.ValueType  # 1
+    """The corresponding dimension is Compressed, only nonzeros are stored."""
+    DIM_SINGLETON: _DimLevelType.ValueType  # 2
+    """The corresponding dimension contains a single coordinate, no sibling
+    elements for each parent.
+    """
+    DIM_LOOSE_COMPRESSED: _DimLevelType.ValueType  # 3
+    """The corresponding dimension is Compressed, but with potential trailing
+    zeros, thus an extra upper bound (high) is used to exclude those zeros.
+    E.g., indices = [1, 2, 0, 0, 3, 4, 0, 0], position = [(0, 2), (4, 6)].
+    """
+
+class DimLevelType(_DimLevelType, metaclass=_DimLevelTypeEnumTypeWrapper):
+    """A DimLevelType indicates the encoding method for a dimension in an array.
+    The semantics of this field are identical to those of the MLIR SparseTensor
+    dialect.
+    This should be kept in sync with the SparseTensor DimLevelType enum:
+    https://github.com/llvm/llvm-project/blob/5674a3c88088e668b684326c2194a6282e8270ff/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td#L86
+    """
+
+DIM_DENSE: DimLevelType.ValueType  # 0
+"""The corresponding dimension is Dense, every entry is stored."""
+DIM_COMPRESSED: DimLevelType.ValueType  # 1
+"""The corresponding dimension is Compressed, only nonzeros are stored."""
+DIM_SINGLETON: DimLevelType.ValueType  # 2
+"""The corresponding dimension contains a single coordinate, no sibling
+elements for each parent.
+"""
+DIM_LOOSE_COMPRESSED: DimLevelType.ValueType  # 3
+"""The corresponding dimension is Compressed, but with potential trailing
+zeros, thus an extra upper bound (high) is used to exclude those zeros.
+E.g., indices = [1, 2, 0, 0, 3, 4, 0, 0], position = [(0, 2), (4, 6)].
+"""
+global___DimLevelType = DimLevelType
+
+class _ProfileType:
+    ValueType = typing.NewType("ValueType", builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+
+class _ProfileTypeEnumTypeWrapper(
+    google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_ProfileType.ValueType], builtins.type
+):
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    INVALID: _ProfileType.ValueType  # 0
+    WINDOW: _ProfileType.ValueType  # 1
+    FLAG: _ProfileType.ValueType  # 2
+    INTEGER: _ProfileType.ValueType  # 3
+
+class ProfileType(_ProfileType, metaclass=_ProfileTypeEnumTypeWrapper):
+    """The type optimization profiles in use for Op-level optimizations."""
+
+INVALID: ProfileType.ValueType  # 0
+WINDOW: ProfileType.ValueType  # 1
+FLAG: ProfileType.ValueType  # 2
+INTEGER: ProfileType.ValueType  # 3
+global___ProfileType = ProfileType
+
+class _ProfileSource:
+    ValueType = typing.NewType("ValueType", builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+
+class _ProfileSourceEnumTypeWrapper(
+    google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_ProfileSource.ValueType], builtins.type
+):
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    PROFILE_SOURCE_UNKNOWN_SOURCE: _ProfileSource.ValueType  # 0
+    PROFILE_SOURCE_EMBEDDED: _ProfileSource.ValueType  # 1
+    PROFILE_SOURCE_REMOTE: _ProfileSource.ValueType  # 2
+
+class ProfileSource(_ProfileSource, metaclass=_ProfileSourceEnumTypeWrapper):
+    """The source of the optimization profile."""
+
+PROFILE_SOURCE_UNKNOWN_SOURCE: ProfileSource.ValueType  # 0
+PROFILE_SOURCE_EMBEDDED: ProfileSource.ValueType  # 1
+PROFILE_SOURCE_REMOTE: ProfileSource.ValueType  # 2
+global___ProfileSource = ProfileSource
+
+class _CompilationEvent:
+    ValueType = typing.NewType("ValueType", builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+
+class _CompilationEventEnumTypeWrapper(
+    google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_CompilationEvent.ValueType], builtins.type
+):
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    COMPILATION_EVENT_UNKNOWN_EVENT: _CompilationEvent.ValueType  # 0
+    COMPILATION_EVENT_FIRST_COMPILATION: _CompilationEvent.ValueType  # 1
+    COMPILATION_EVENT_RECOMPILATION: _CompilationEvent.ValueType  # 2
+
+class CompilationEvent(_CompilationEvent, metaclass=_CompilationEventEnumTypeWrapper):
+    """The compilation event that triggered the use of the profile."""
+
+COMPILATION_EVENT_UNKNOWN_EVENT: CompilationEvent.ValueType  # 0
+COMPILATION_EVENT_FIRST_COMPILATION: CompilationEvent.ValueType  # 1
+COMPILATION_EVENT_RECOMPILATION: CompilationEvent.ValueType  # 2
+global___CompilationEvent = CompilationEvent
+
+class _PaddingType:
+    ValueType = typing.NewType("ValueType", builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+
+class _PaddingTypeEnumTypeWrapper(
+    google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_PaddingType.ValueType], builtins.type
+):
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    PADDING_INVALID: _PaddingType.ValueType  # 0
+    PADDING_VALID: _PaddingType.ValueType  # 1
+    """Only valid portion of the base are covered."""
+    PADDING_SAME: _PaddingType.ValueType  # 2
+    """Extra is added to produce same output size as the input."""
+
+class PaddingType(_PaddingType, metaclass=_PaddingTypeEnumTypeWrapper): ...
+
+PADDING_INVALID: PaddingType.ValueType  # 0
+PADDING_VALID: PaddingType.ValueType  # 1
+"""Only valid portion of the base are covered."""
+PADDING_SAME: PaddingType.ValueType  # 2
+"""Extra is added to produce same output size as the input."""
+global___PaddingType = PaddingType
+
+class _FftType:
+    ValueType = typing.NewType("ValueType", builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+
+class _FftTypeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_FftType.ValueType], builtins.type):
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    FFT: _FftType.ValueType  # 0
+    """Forward FFT; complex in, complex out."""
+    IFFT: _FftType.ValueType  # 1
+    """Inverse FFT; complex in, complex out."""
+    RFFT: _FftType.ValueType  # 2
+    """Forward real FFT; real in, fft_length / 2 + 1 complex out"""
+    IRFFT: _FftType.ValueType  # 3
+    """Inverse real FFT; fft_length / 2 + 1 complex in,"""
+
+class FftType(_FftType, metaclass=_FftTypeEnumTypeWrapper): ...
+
+FFT: FftType.ValueType  # 0
+"""Forward FFT; complex in, complex out."""
+IFFT: FftType.ValueType  # 1
+"""Inverse FFT; complex in, complex out."""
+RFFT: FftType.ValueType  # 2
+"""Forward real FFT; real in, fft_length / 2 + 1 complex out"""
+IRFFT: FftType.ValueType  # 3
+"""Inverse real FFT; fft_length / 2 + 1 complex in,"""
+global___FftType = FftType
+
+class _SparsityType:
+    ValueType = typing.NewType("ValueType", builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+
+class _SparsityTypeEnumTypeWrapper(
+    google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_SparsityType.ValueType], builtins.type
+):
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    SPARSITY_INVALID: _SparsityType.ValueType  # 0
+    SPARSITY_STRUCTURED_N_M: _SparsityType.ValueType  # 1
+    """Structured N:M sparsity."""
+
+class SparsityType(_SparsityType, metaclass=_SparsityTypeEnumTypeWrapper): ...
+
+SPARSITY_INVALID: SparsityType.ValueType  # 0
+SPARSITY_STRUCTURED_N_M: SparsityType.ValueType  # 1
+"""Structured N:M sparsity."""
+global___SparsityType = SparsityType
+
+class _RandomDistribution:
+    ValueType = typing.NewType("ValueType", builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+
+class _RandomDistributionEnumTypeWrapper(
+    google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_RandomDistribution.ValueType], builtins.type
+):
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    RNG_INVALID: _RandomDistribution.ValueType  # 0
+    RNG_UNIFORM: _RandomDistribution.ValueType  # 1
+    """Creates a uniform-distribution-generated random number on the semi-open
+    interval [parameter[0], parameter[1]).
+    """
+    RNG_NORMAL: _RandomDistribution.ValueType  # 2
+    """Creates a normal-distribution-generated random number with mean
+    parameter[0] and standard deviation parameter[1].
+    """
+
+class RandomDistribution(_RandomDistribution, metaclass=_RandomDistributionEnumTypeWrapper): ...
+
+RNG_INVALID: RandomDistribution.ValueType  # 0
+RNG_UNIFORM: RandomDistribution.ValueType  # 1
+"""Creates a uniform-distribution-generated random number on the semi-open
+interval [parameter[0], parameter[1]).
+"""
+RNG_NORMAL: RandomDistribution.ValueType  # 2
+"""Creates a normal-distribution-generated random number with mean
+parameter[0] and standard deviation parameter[1].
+"""
+global___RandomDistribution = RandomDistribution
+
+class _RandomAlgorithm:
+    ValueType = typing.NewType("ValueType", builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+
+class _RandomAlgorithmEnumTypeWrapper(
+    google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_RandomAlgorithm.ValueType], builtins.type
+):
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    RNG_DEFAULT: _RandomAlgorithm.ValueType  # 0
+    """Backend dependent default algorithm."""
+    RNG_THREE_FRY: _RandomAlgorithm.ValueType  # 1
+    RNG_PHILOX: _RandomAlgorithm.ValueType  # 2
+    """Next: 2"""
+
+class RandomAlgorithm(_RandomAlgorithm, metaclass=_RandomAlgorithmEnumTypeWrapper): ...
+
+RNG_DEFAULT: RandomAlgorithm.ValueType  # 0
+"""Backend dependent default algorithm."""
+RNG_THREE_FRY: RandomAlgorithm.ValueType  # 1
+RNG_PHILOX: RandomAlgorithm.ValueType  # 2
+"""Next: 2"""
+global___RandomAlgorithm = RandomAlgorithm
+
+@typing.final
+class PaddingConfig(google.protobuf.message.Message):
+    """Describes the padding configuration for Pad operation. The padding amount on
+    both edges as well as between the elements are specified for each dimension.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.final
+    class PaddingConfigDimension(google.protobuf.message.Message):
+        """Describes the padding configuration for a dimension."""
+
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        EDGE_PADDING_LOW_FIELD_NUMBER: builtins.int
+        EDGE_PADDING_HIGH_FIELD_NUMBER: builtins.int
+        INTERIOR_PADDING_FIELD_NUMBER: builtins.int
+        edge_padding_low: builtins.int
+        """Padding amount on the low-end (next to the index 0). May be negative."""
+        edge_padding_high: builtins.int
+        """Padding amount on the high-end (next to the highest index). May be
+        negative.
+        """
+        interior_padding: builtins.int
+        """Padding amount between the elements. May not be negative."""
+        def __init__(
+            self,
+            *,
+            edge_padding_low: builtins.int | None = ...,
+            edge_padding_high: builtins.int | None = ...,
+            interior_padding: builtins.int | None = ...,
+        ) -> None: ...
+        def ClearField(
+            self,
+            field_name: typing.Literal[
+                "edge_padding_high",
+                b"edge_padding_high",
+                "edge_padding_low",
+                b"edge_padding_low",
+                "interior_padding",
+                b"interior_padding",
+            ],
+        ) -> None: ...
+
+    DIMENSIONS_FIELD_NUMBER: builtins.int
+    @property
+    def dimensions(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___PaddingConfig.PaddingConfigDimension]:
+        """The padding configuration for all dimensions."""
+
+    def __init__(
+        self, *, dimensions: collections.abc.Iterable[global___PaddingConfig.PaddingConfigDimension] | None = ...
+    ) -> None: ...
+    def ClearField(self, field_name: typing.Literal["dimensions", b"dimensions"]) -> None: ...
+
+global___PaddingConfig = PaddingConfig
+
+@typing.final
+class TileProto(google.protobuf.message.Message):
+    """Describes a tile used in tiling-based layout. Refer to
+    g3doc/third_party/xla/docs/tiled_layout.md for details about tiling-based
+    layout.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    DIMENSIONS_FIELD_NUMBER: builtins.int
+    @property
+    def dimensions(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """Number of elements in each dimension of the tile. It's ordered from the
+        most major dimension of the tile to the most minor dimension of the tile.
+        The dimensions correspond to a suffix of the dimensions of the shape being
+        tiled.
+        """
+
+    def __init__(self, *, dimensions: collections.abc.Iterable[builtins.int] | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["dimensions", b"dimensions"]) -> None: ...
+
+global___TileProto = TileProto
+
+@typing.final
+class SplitConfigProto(google.protobuf.message.Message):
+    """Describes how data should be split between different memories."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    DIMENSION_FIELD_NUMBER: builtins.int
+    SPLIT_INDICES_FIELD_NUMBER: builtins.int
+    dimension: builtins.int
+    """The dimension that is split."""
+    @property
+    def split_indices(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The indices where each split point occurs. For example, if the dimension
+        size is 1024, a split_indices value of {512} indicates a two-way split of
+        data through the middle.
+        """
+
+    def __init__(
+        self, *, dimension: builtins.int | None = ..., split_indices: collections.abc.Iterable[builtins.int] | None = ...
+    ) -> None: ...
+    def ClearField(self, field_name: typing.Literal["dimension", b"dimension", "split_indices", b"split_indices"]) -> None: ...
+
+global___SplitConfigProto = SplitConfigProto
+
+@typing.final
+class LayoutProto(google.protobuf.message.Message):
+    """A layout describes how the array is placed in (1D) memory space.  This
+    includes the minor-to-major ordering of dimensions within a shape.
+
+    Clients must specify the layouts of input Literals to the
+    computation. Layouts specified in interior operations which take Shapes (for
+    example, Convert) are ignored.
+
+    See the XLA documentation for more information on shapes and layouts.
+
+    LINT.IfChange
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    DIM_LEVEL_TYPES_FIELD_NUMBER: builtins.int
+    DIM_UNIQUE_FIELD_NUMBER: builtins.int
+    DIM_ORDERED_FIELD_NUMBER: builtins.int
+    MINOR_TO_MAJOR_FIELD_NUMBER: builtins.int
+    TILES_FIELD_NUMBER: builtins.int
+    TAIL_PADDING_ALIGNMENT_IN_ELEMENTS_FIELD_NUMBER: builtins.int
+    ELEMENT_SIZE_IN_BITS_FIELD_NUMBER: builtins.int
+    MEMORY_SPACE_FIELD_NUMBER: builtins.int
+    INDEX_PRIMITIVE_TYPE_FIELD_NUMBER: builtins.int
+    POINTER_PRIMITIVE_TYPE_FIELD_NUMBER: builtins.int
+    PHYSICAL_SHAPE_FIELD_NUMBER: builtins.int
+    DYNAMIC_SHAPE_METADATA_PREFIX_BYTES_FIELD_NUMBER: builtins.int
+    SPLIT_CONFIGS_FIELD_NUMBER: builtins.int
+    tail_padding_alignment_in_elements: builtins.int
+    """The shape is padded at the end to multiple of, in terms of number of
+    elements. This is useful when tiling does not bring the shape to certain
+    desired granules. Tiling effectively pads/reshapes/transposes the shape
+    to another shape. This field pads the total number of elements of that
+    new shape to a multiple of certain number of elements. This is useful such
+    as we want a layout which does not tile the data but still requires it to
+    be padded to certain number of elements.
+    """
+    element_size_in_bits: builtins.int
+    """(Optional) Bit size of each element. When unspecified or being 0, default
+    to ShapeUtil::ByteSizeOfPrimitiveType.
+    """
+    memory_space: builtins.int
+    """Memory space where this array resides. The integer field is interpreted in
+    a backend-specific manner.
+    """
+    index_primitive_type: global___PrimitiveType.ValueType
+    """The integer types to be used for indices and pointers.  These fields must
+    not be used unless the layout represents a sparse array.  The PrimitiveType
+    must correspond to an unsigned integer (U8, U16, U32, or U64).
+    If not provided, the compiler will use the largest unsigned integer
+    that is naturally supported by the target device (U32 or U64 in currently
+    supported devices).
+    """
+    pointer_primitive_type: global___PrimitiveType.ValueType
+    dynamic_shape_metadata_prefix_bytes: builtins.int
+    """The dynamic shape metadata size in bytes in front of the shape data. The
+    field may be non-zero for a static shape whose associated buffer is for a
+    dynamic shape, e.g. a result of SliceToDynamic.
+    """
+    @property
+    def dim_level_types(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[global___DimLevelType.ValueType]:
+        """The dimension level type list for this array, specifying the way in which
+        each array dimension is represented in memory. If this list is empty, the
+        array is assumed to be dense.
+        """
+
+    @property
+    def dim_unique(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.bool]:
+        """Whether each dimension is unique or ordered.  Each of the following lists
+        must be empty, or have one entry for each entry of dim_level_types.  If
+        either list is empty, all dimensions are assumed to be unique and ordered,
+        respectively.  Entries in this list may not be false for some DimLevelType
+        values (such as DIM_DENSE in particular).
+        """
+
+    @property
+    def dim_ordered(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.bool]: ...
+    @property
+    def minor_to_major(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """Sequence of dimension numbers, from minor (fastest varying index) to major
+        (slowest varying index). This field is required.
+        """
+
+    @property
+    def tiles(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___TileProto]:
+        """A sequence of tiles, starting from the tile that's applied first to the
+        Shape.
+
+        TODO(b/119839262): implement tiling in each backend or add Unimplemented
+        error.
+        """
+
+    @property
+    def physical_shape(self) -> global___ShapeProto:
+        """The physical, on-device shape used to represent the shape this layout
+        belongs to. Only used for sparse arrays.
+        The layout(s) contained within the physical shape should not also contain
+        a physical shape.
+        """
+
+    @property
+    def split_configs(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___SplitConfigProto]:
+        """The split configurations which describe if/how the data is split between
+        different memories.
+        """
+
+    def __init__(
+        self,
+        *,
+        dim_level_types: collections.abc.Iterable[global___DimLevelType.ValueType] | None = ...,
+        dim_unique: collections.abc.Iterable[builtins.bool] | None = ...,
+        dim_ordered: collections.abc.Iterable[builtins.bool] | None = ...,
+        minor_to_major: collections.abc.Iterable[builtins.int] | None = ...,
+        tiles: collections.abc.Iterable[global___TileProto] | None = ...,
+        tail_padding_alignment_in_elements: builtins.int | None = ...,
+        element_size_in_bits: builtins.int | None = ...,
+        memory_space: builtins.int | None = ...,
+        index_primitive_type: global___PrimitiveType.ValueType | None = ...,
+        pointer_primitive_type: global___PrimitiveType.ValueType | None = ...,
+        physical_shape: global___ShapeProto | None = ...,
+        dynamic_shape_metadata_prefix_bytes: builtins.int | None = ...,
+        split_configs: collections.abc.Iterable[global___SplitConfigProto] | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["physical_shape", b"physical_shape"]) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "dim_level_types",
+            b"dim_level_types",
+            "dim_ordered",
+            b"dim_ordered",
+            "dim_unique",
+            b"dim_unique",
+            "dynamic_shape_metadata_prefix_bytes",
+            b"dynamic_shape_metadata_prefix_bytes",
+            "element_size_in_bits",
+            b"element_size_in_bits",
+            "index_primitive_type",
+            b"index_primitive_type",
+            "memory_space",
+            b"memory_space",
+            "minor_to_major",
+            b"minor_to_major",
+            "physical_shape",
+            b"physical_shape",
+            "pointer_primitive_type",
+            b"pointer_primitive_type",
+            "split_configs",
+            b"split_configs",
+            "tail_padding_alignment_in_elements",
+            b"tail_padding_alignment_in_elements",
+            "tiles",
+            b"tiles",
+        ],
+    ) -> None: ...
+
+global___LayoutProto = LayoutProto
+
+@typing.final
+class ShapeProto(google.protobuf.message.Message):
+    """A shape describes the number of dimensions in the array, the size of each
+    dimension, and the primitive component type.
+
+    Tuples are a special case in that they have rank zero and have tuple_shapes
+    defined.
+
+    See the XLA documentation for more information on shapes and layouts.
+
+    LINT.IfChange
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    ELEMENT_TYPE_FIELD_NUMBER: builtins.int
+    DIMENSIONS_FIELD_NUMBER: builtins.int
+    TUPLE_SHAPES_FIELD_NUMBER: builtins.int
+    LAYOUT_FIELD_NUMBER: builtins.int
+    IS_DYNAMIC_DIMENSION_FIELD_NUMBER: builtins.int
+    element_type: global___PrimitiveType.ValueType
+    """The element type for this shape."""
+    @property
+    def dimensions(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The size (number of elements) for each dimension, or an upper bound on the
+        size if the dimension is dynamic.  In XLA, dimensions are numbered from 0
+        to N-1 for an N-dimensional array. The first element of 'dimensions' is the
+        size of dimension 0, the second element is the size of dimension 1, and so
+        forth.  Empty list indicates a scalar.
+
+        If the respective element in 'is_dimension_dynamic' is true then the value
+        in this field represents an upper bound on the size of the dimension.
+        """
+
+    @property
+    def tuple_shapes(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ShapeProto]:
+        """For tuples only, the shapes of constituent shapes in the tuple sequence."""
+
+    @property
+    def layout(self) -> global___LayoutProto:
+        """The layout used to back this shape."""
+
+    @property
+    def is_dynamic_dimension(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.bool]:
+        """For arrays, this indicates whether or not each dimension is
+        dynamically-sized. The number of elements in this repeated field should be
+        zero (indicating that no dimensions are dynamic) or equal to the number of
+        elements in the 'dimensions' field.
+        """
+
+    def __init__(
+        self,
+        *,
+        element_type: global___PrimitiveType.ValueType | None = ...,
+        dimensions: collections.abc.Iterable[builtins.int] | None = ...,
+        tuple_shapes: collections.abc.Iterable[global___ShapeProto] | None = ...,
+        layout: global___LayoutProto | None = ...,
+        is_dynamic_dimension: collections.abc.Iterable[builtins.bool] | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["layout", b"layout"]) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "dimensions",
+            b"dimensions",
+            "element_type",
+            b"element_type",
+            "is_dynamic_dimension",
+            b"is_dynamic_dimension",
+            "layout",
+            b"layout",
+            "tuple_shapes",
+            b"tuple_shapes",
+        ],
+    ) -> None: ...
+
+global___ShapeProto = ShapeProto
+
+@typing.final
+class ProgramShapeProto(google.protobuf.message.Message):
+    """Shape of the parameters and output of a computation (like a traditional
+    function signature).
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    PARAMETERS_FIELD_NUMBER: builtins.int
+    RESULT_FIELD_NUMBER: builtins.int
+    PARAMETER_NAMES_FIELD_NUMBER: builtins.int
+    @property
+    def parameters(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ShapeProto]: ...
+    @property
+    def result(self) -> global___ShapeProto: ...
+    @property
+    def parameter_names(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: ...
+    def __init__(
+        self,
+        *,
+        parameters: collections.abc.Iterable[global___ShapeProto] | None = ...,
+        result: global___ShapeProto | None = ...,
+        parameter_names: collections.abc.Iterable[builtins.str] | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["result", b"result"]) -> builtins.bool: ...
+    def ClearField(
+        self, field_name: typing.Literal["parameter_names", b"parameter_names", "parameters", b"parameters", "result", b"result"]
+    ) -> None: ...
+
+global___ProgramShapeProto = ProgramShapeProto
+
+@typing.final
+class ComputationStats(google.protobuf.message.Message):
+    """Statistics of a computation."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    FLOP_COUNT_FIELD_NUMBER: builtins.int
+    TRANSCENDENTAL_COUNT_FIELD_NUMBER: builtins.int
+    flop_count: builtins.float
+    """The number of floating point operations in the computation."""
+    transcendental_count: builtins.float
+    """The number of transcendental operations (e.g., exp) in the computation."""
+    def __init__(self, *, flop_count: builtins.float | None = ..., transcendental_count: builtins.float | None = ...) -> None: ...
+    def ClearField(
+        self, field_name: typing.Literal["flop_count", b"flop_count", "transcendental_count", b"transcendental_count"]
+    ) -> None: ...
+
+global___ComputationStats = ComputationStats
+
+@typing.final
+class OpMetadata(google.protobuf.message.Message):
+    """Symbolization metadata for HLO Instructions.
+
+    This metadata is used for debugging XLA code generation, as well as
+    performance profiling of XLA-generated executables.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.final
+    class ProfileInfo(google.protobuf.message.Message):
+        """Information about the optimization profile that this operation 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
+        relative_speedup: builtins.float
+        """Speedup of tuned config compared to default config.
+        TODO(b/203817882) Set the relative_speedup.
+        """
+        profile_source: global___ProfileSource.ValueType
+        """The source of the optimization profiles that this operation contains."""
+        compilation_event: global___CompilationEvent.ValueType
+        """The compilation event that triggered the use of the profiles."""
+        @property
+        def profile_type(
+            self,
+        ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[global___ProfileType.ValueType]:
+            """The type of optimization profiles that this operation contains."""
+
+        def __init__(
+            self,
+            *,
+            profile_type: collections.abc.Iterable[global___ProfileType.ValueType] | None = ...,
+            relative_speedup: builtins.float | None = ...,
+            profile_source: global___ProfileSource.ValueType | None = ...,
+            compilation_event: global___CompilationEvent.ValueType | None = ...,
+        ) -> None: ...
+        def ClearField(
+            self,
+            field_name: typing.Literal[
+                "compilation_event",
+                b"compilation_event",
+                "profile_source",
+                b"profile_source",
+                "profile_type",
+                b"profile_type",
+                "relative_speedup",
+                b"relative_speedup",
+            ],
+        ) -> None: ...
+
+    OP_TYPE_FIELD_NUMBER: builtins.int
+    OP_NAME_FIELD_NUMBER: builtins.int
+    SOURCE_FILE_FIELD_NUMBER: builtins.int
+    SOURCE_LINE_FIELD_NUMBER: builtins.int
+    PROFILE_TYPE_FIELD_NUMBER: builtins.int
+    SIZE_OF_GENERATED_CODE_IN_BYTES_FIELD_NUMBER: builtins.int
+    SIZE_OF_MEMORY_WORKING_SET_IN_BYTES_FIELD_NUMBER: builtins.int
+    PROFILE_INFO_FIELD_NUMBER: builtins.int
+    DEDUPLICATED_NAME_FIELD_NUMBER: builtins.int
+    PRESERVE_LAYOUT_FIELD_NUMBER: builtins.int
+    STACK_FRAME_ID_FIELD_NUMBER: builtins.int
+    SCHEDULING_NAME_FIELD_NUMBER: builtins.int
+    op_type: builtins.str
+    """The framework op name that generated this XLA op.
+
+    Frameworks that build on top of XLA should mirror the names of their ops
+    back to users by specifying the op_type. In this way, even if the
+    framework's "ops" are implemented as multiple XLA HLO Ops, they can be
+    grouped appropriately. (e.g. if a SoftMax layer is emitted into XLA as
+    multiple ops, then each op should have the op_type be "SoftMax".)
+    """
+    op_name: builtins.str
+    """The user-specified name of the op.
+
+    This name is often unique within a computation. Note: some frameworks
+    add auto-generated names if the user does not provide one.
+    """
+    source_file: builtins.str
+    """Indicate a file and line that this op is associated to in a user's program.
+
+    e.g. it could be the file and line of user code that generated the op.
+    """
+    source_line: builtins.int
+    size_of_generated_code_in_bytes: builtins.int
+    """The footprint of the generated code for the instruction."""
+    size_of_memory_working_set_in_bytes: builtins.int
+    """The size of the working set, i.e., the amount of memory, used by the
+    instruction in a compiler-managed fast device memory.
+    """
+    deduplicated_name: builtins.str
+    """Deduplicated HLO name for this op. In some cases, we can have multiple
+    instructions (e.g. fusions) that are considered duplicates. We want to
+    group them together under the same name so that we can group them together
+    during analysis (e.g. HLO Op Profile tool in Xprof).
+    E.g. If we have fusion.1, fusion.2, and fusion.3 marked as duplicates,
+    fusion.2 and fusion.3 will have deduplicated_name = fusion.1
+    """
+    preserve_layout: builtins.bool
+    """Whether to preserve the layout of the HLO op."""
+    stack_frame_id: builtins.int
+    """1-based position of the frame in frames flat array.
+    Ids are 1-based to keep 0 value as representation of non-set property.
+    """
+    scheduling_name: builtins.str
+    """Instruction name available upon scheduling."""
+    @property
+    def profile_type(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[global___ProfileType.ValueType]:
+        """Deprecated, use [ProfileInfo][profile_type] instead."""
+
+    @property
+    def profile_info(self) -> global___OpMetadata.ProfileInfo:
+        """Profile information for the Op."""
+
+    def __init__(
+        self,
+        *,
+        op_type: builtins.str | None = ...,
+        op_name: builtins.str | None = ...,
+        source_file: builtins.str | None = ...,
+        source_line: builtins.int | None = ...,
+        profile_type: collections.abc.Iterable[global___ProfileType.ValueType] | None = ...,
+        size_of_generated_code_in_bytes: builtins.int | None = ...,
+        size_of_memory_working_set_in_bytes: builtins.int | None = ...,
+        profile_info: global___OpMetadata.ProfileInfo | None = ...,
+        deduplicated_name: builtins.str | None = ...,
+        preserve_layout: builtins.bool | None = ...,
+        stack_frame_id: builtins.int | None = ...,
+        scheduling_name: builtins.str | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["profile_info", b"profile_info"]) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "deduplicated_name",
+            b"deduplicated_name",
+            "op_name",
+            b"op_name",
+            "op_type",
+            b"op_type",
+            "preserve_layout",
+            b"preserve_layout",
+            "profile_info",
+            b"profile_info",
+            "profile_type",
+            b"profile_type",
+            "scheduling_name",
+            b"scheduling_name",
+            "size_of_generated_code_in_bytes",
+            b"size_of_generated_code_in_bytes",
+            "size_of_memory_working_set_in_bytes",
+            b"size_of_memory_working_set_in_bytes",
+            "source_file",
+            b"source_file",
+            "source_line",
+            b"source_line",
+            "stack_frame_id",
+            b"stack_frame_id",
+        ],
+    ) -> None: ...
+
+global___OpMetadata = OpMetadata
+
+@typing.final
+class ExecutionProfile(google.protobuf.message.Message):
+    """Profile data from the execution of a computation."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    COMPILATION_CACHE_HIT_FIELD_NUMBER: builtins.int
+    COMPILE_TIME_MS_FIELD_NUMBER: builtins.int
+    COMPUTE_CYCLE_COUNT_FIELD_NUMBER: builtins.int
+    COMPUTE_TIME_NS_FIELD_NUMBER: builtins.int
+    COMPUTE_AND_TRANSFER_TIME_NS_FIELD_NUMBER: builtins.int
+    EXECUTABLE_SIZE_IN_BYTES_FIELD_NUMBER: builtins.int
+    PROFILE_CACHE_HIT_FIELD_NUMBER: builtins.int
+    WARMUP_RUN_EXECUTED_FIELD_NUMBER: builtins.int
+    compilation_cache_hit: builtins.bool
+    """Whether the executable was read from the compilation cache."""
+    compile_time_ms: builtins.int
+    """The time in milliseconds spent to compile the computation. This only set if
+    the executable was not read from the compilation cache
+    (compilation_cache_hit == false).
+    """
+    compute_cycle_count: builtins.int
+    """The number of cycles spent for the computation. This does not include the
+    time taken for the data transfers between the host and the device. This is
+    a target-dependent field and only used for debugging purposes.
+    """
+    compute_time_ns: builtins.int
+    """The time in nanoseconds spent for the computation, without data transfer."""
+    compute_and_transfer_time_ns: builtins.int
+    """The time in nanoseconds spent for the entire computation, including the
+    result data transfer time. Current implementation does not spend any cycles
+    for the input data transfer since the memory is initialized with the proper
+    values before the execution.
+    """
+    executable_size_in_bytes: builtins.int
+    """The size of the binary code in the executable."""
+    profile_cache_hit: builtins.bool
+    """Whether this profile was drawn from a cache of profiles instead of from
+    execution on the hardware.
+    """
+    warmup_run_executed: builtins.bool
+    """Whether a warm-up run of the computation was executed before the
+    measured execution.
+    """
+    def __init__(
+        self,
+        *,
+        compilation_cache_hit: builtins.bool | None = ...,
+        compile_time_ms: builtins.int | None = ...,
+        compute_cycle_count: builtins.int | None = ...,
+        compute_time_ns: builtins.int | None = ...,
+        compute_and_transfer_time_ns: builtins.int | None = ...,
+        executable_size_in_bytes: builtins.int | None = ...,
+        profile_cache_hit: builtins.bool | None = ...,
+        warmup_run_executed: builtins.bool | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "compilation_cache_hit",
+            b"compilation_cache_hit",
+            "compile_time_ms",
+            b"compile_time_ms",
+            "compute_and_transfer_time_ns",
+            b"compute_and_transfer_time_ns",
+            "compute_cycle_count",
+            b"compute_cycle_count",
+            "compute_time_ns",
+            b"compute_time_ns",
+            "executable_size_in_bytes",
+            b"executable_size_in_bytes",
+            "profile_cache_hit",
+            b"profile_cache_hit",
+            "warmup_run_executed",
+            b"warmup_run_executed",
+        ],
+    ) -> None: ...
+
+global___ExecutionProfile = ExecutionProfile
+
+@typing.final
+class ExecutionHandle(google.protobuf.message.Message):
+    """Handle given to a user that represents an execution that the user launched
+    asynchronously on the device.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    HANDLE_FIELD_NUMBER: builtins.int
+    handle: builtins.int
+    def __init__(self, *, handle: builtins.int | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["handle", b"handle"]) -> None: ...
+
+global___ExecutionHandle = ExecutionHandle
+
+@typing.final
+class GlobalDataHandle(google.protobuf.message.Message):
+    """Handle given to a user that represents a globally accessible allocation.
+    Contrast this against a ComputationDataHandle, which is not globally
+    accessible, since it only exists within a specific computation.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    HANDLE_FIELD_NUMBER: builtins.int
+    handle: builtins.int
+    def __init__(self, *, handle: builtins.int | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["handle", b"handle"]) -> None: ...
+
+global___GlobalDataHandle = GlobalDataHandle
+
+@typing.final
+class DeviceHandle(google.protobuf.message.Message):
+    """Handle given to a user that represents a replicated virtual device. Each
+    replicated device represents N physical devices for execution where N is the
+    number of replicas.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    HANDLE_FIELD_NUMBER: builtins.int
+    DEVICE_COUNT_FIELD_NUMBER: builtins.int
+    handle: builtins.int
+    device_count: builtins.int
+    """The number of model-parallel virtual devices that communicate via XLA
+    Send/Recv instructions.
+    """
+    def __init__(self, *, handle: builtins.int | None = ..., device_count: builtins.int | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["device_count", b"device_count", "handle", b"handle"]) -> None: ...
+
+global___DeviceHandle = DeviceHandle
+
+@typing.final
+class ChannelHandle(google.protobuf.message.Message):
+    """Handle given to a user to represent a channel between two computations
+    via a Send and Recv instruction pair. Channels are unbuffered, so Send
+    Send instructions will be blocked until the data is transferred.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    class _ChannelType:
+        ValueType = typing.NewType("ValueType", builtins.int)
+        V: typing_extensions.TypeAlias = ValueType
+
+    class _ChannelTypeEnumTypeWrapper(
+        google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ChannelHandle._ChannelType.ValueType], builtins.type
+    ):
+        DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+        CHANNEL_TYPE_INVALID: ChannelHandle._ChannelType.ValueType  # 0
+        """Invalid primitive type to serve as default."""
+        DEVICE_TO_DEVICE: ChannelHandle._ChannelType.ValueType  # 1
+        """A channel for sending data between devices."""
+        DEVICE_TO_HOST: ChannelHandle._ChannelType.ValueType  # 2
+        """A channel for sending data from the device to the host. Can only be used
+        with a Send operation.
+        """
+        HOST_TO_DEVICE: ChannelHandle._ChannelType.ValueType  # 3
+        """A channel for sending data from the host to the device. Can only be used
+        with a Recv operation.
+        """
+
+    class ChannelType(_ChannelType, metaclass=_ChannelTypeEnumTypeWrapper): ...
+    CHANNEL_TYPE_INVALID: ChannelHandle.ChannelType.ValueType  # 0
+    """Invalid primitive type to serve as default."""
+    DEVICE_TO_DEVICE: ChannelHandle.ChannelType.ValueType  # 1
+    """A channel for sending data between devices."""
+    DEVICE_TO_HOST: ChannelHandle.ChannelType.ValueType  # 2
+    """A channel for sending data from the device to the host. Can only be used
+    with a Send operation.
+    """
+    HOST_TO_DEVICE: ChannelHandle.ChannelType.ValueType  # 3
+    """A channel for sending data from the host to the device. Can only be used
+    with a Recv operation.
+    """
+
+    HANDLE_FIELD_NUMBER: builtins.int
+    TYPE_FIELD_NUMBER: builtins.int
+    handle: builtins.int
+    type: global___ChannelHandle.ChannelType.ValueType
+    def __init__(
+        self, *, handle: builtins.int | None = ..., type: global___ChannelHandle.ChannelType.ValueType | None = ...
+    ) -> None: ...
+    def ClearField(self, field_name: typing.Literal["handle", b"handle", "type", b"type"]) -> None: ...
+
+global___ChannelHandle = ChannelHandle
+
+@typing.final
+class DeviceAssignmentProto(google.protobuf.message.Message):
+    """DeviceAssignmentProto is a serialized form of DeviceAssignment class, which
+    represents the device ids assigned to a set of replicated computations.
+    See xla::DeviceAssignment class comment for more details.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.final
+    class ComputationDevice(google.protobuf.message.Message):
+        """Each logical computation runs on replica_count physical devices.
+        ComputationDevice represents the device ids assinged to the replicas.
+        """
+
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        REPLICA_DEVICE_IDS_FIELD_NUMBER: builtins.int
+        @property
+        def replica_device_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ...
+        def __init__(self, *, replica_device_ids: collections.abc.Iterable[builtins.int] | None = ...) -> None: ...
+        def ClearField(self, field_name: typing.Literal["replica_device_ids", b"replica_device_ids"]) -> None: ...
+
+    REPLICA_COUNT_FIELD_NUMBER: builtins.int
+    COMPUTATION_COUNT_FIELD_NUMBER: builtins.int
+    COMPUTATION_DEVICES_FIELD_NUMBER: builtins.int
+    replica_count: builtins.int
+    computation_count: builtins.int
+    @property
+    def computation_devices(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[
+        global___DeviceAssignmentProto.ComputationDevice
+    ]: ...
+    def __init__(
+        self,
+        *,
+        replica_count: builtins.int | None = ...,
+        computation_count: builtins.int | None = ...,
+        computation_devices: collections.abc.Iterable[global___DeviceAssignmentProto.ComputationDevice] | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "computation_count",
+            b"computation_count",
+            "computation_devices",
+            b"computation_devices",
+            "replica_count",
+            b"replica_count",
+        ],
+    ) -> None: ...
+
+global___DeviceAssignmentProto = DeviceAssignmentProto
+
+@typing.final
+class LiteralProto(google.protobuf.message.Message):
+    """Literals are used when the server and client need to exchange materialized
+    data / results. Literals are also used to describe constants used in
+    computations.
+
+    Transfers to/from the client are encoded in literal form, and the structure
+    of the repeated fields is implied by the shape.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    SHAPE_FIELD_NUMBER: builtins.int
+    PREDS_FIELD_NUMBER: builtins.int
+    S2S_FIELD_NUMBER: builtins.int
+    S4S_FIELD_NUMBER: builtins.int
+    S8S_FIELD_NUMBER: builtins.int
+    U2S_FIELD_NUMBER: builtins.int
+    U4S_FIELD_NUMBER: builtins.int
+    U8S_FIELD_NUMBER: builtins.int
+    S32S_FIELD_NUMBER: builtins.int
+    S64S_FIELD_NUMBER: builtins.int
+    U32S_FIELD_NUMBER: builtins.int
+    U64S_FIELD_NUMBER: builtins.int
+    F32S_FIELD_NUMBER: builtins.int
+    F64S_FIELD_NUMBER: builtins.int
+    C64S_FIELD_NUMBER: builtins.int
+    C128S_FIELD_NUMBER: builtins.int
+    TUPLE_LITERALS_FIELD_NUMBER: builtins.int
+    F16S_FIELD_NUMBER: builtins.int
+    BF16S_FIELD_NUMBER: builtins.int
+    U16S_FIELD_NUMBER: builtins.int
+    S16S_FIELD_NUMBER: builtins.int
+    F8E5M2S_FIELD_NUMBER: builtins.int
+    F8E4M3FNS_FIELD_NUMBER: builtins.int
+    F8E4M3B11FNUZS_FIELD_NUMBER: builtins.int
+    F8E5M2FNUZS_FIELD_NUMBER: builtins.int
+    F8E4M3FNUZS_FIELD_NUMBER: builtins.int
+    SPARSE_INDICES_FIELD_NUMBER: builtins.int
+    s2s: builtins.bytes
+    s4s: builtins.bytes
+    s8s: builtins.bytes
+    u2s: builtins.bytes
+    u4s: builtins.bytes
+    u8s: builtins.bytes
+    f16s: builtins.bytes
+    """The F16s, BF16s, U16s and S16s are encoded in little endian byte order"""
+    bf16s: builtins.bytes
+    u16s: builtins.bytes
+    s16s: builtins.bytes
+    f8e5m2s: builtins.bytes
+    f8e4m3fns: builtins.bytes
+    f8e4m3b11fnuzs: builtins.bytes
+    f8e5m2fnuzs: builtins.bytes
+    f8e4m3fnuzs: builtins.bytes
+    @property
+    def shape(self) -> global___ShapeProto: ...
+    @property
+    def preds(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.bool]: ...
+    @property
+    def s32s(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ...
+    @property
+    def s64s(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ...
+    @property
+    def u32s(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ...
+    @property
+    def u64s(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ...
+    @property
+    def f32s(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.float]: ...
+    @property
+    def f64s(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.float]: ...
+    @property
+    def c64s(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.float]:
+        """Stored as interleaved real, imag floats."""
+
+    @property
+    def c128s(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.float]:
+        """Stored as interleaved real, imag doubles."""
+
+    @property
+    def tuple_literals(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___LiteralProto]: ...
+    @property
+    def sparse_indices(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """Next = 28"""
+
+    def __init__(
+        self,
+        *,
+        shape: global___ShapeProto | None = ...,
+        preds: collections.abc.Iterable[builtins.bool] | None = ...,
+        s2s: builtins.bytes | None = ...,
+        s4s: builtins.bytes | None = ...,
+        s8s: builtins.bytes | None = ...,
+        u2s: builtins.bytes | None = ...,
+        u4s: builtins.bytes | None = ...,
+        u8s: builtins.bytes | None = ...,
+        s32s: collections.abc.Iterable[builtins.int] | None = ...,
+        s64s: collections.abc.Iterable[builtins.int] | None = ...,
+        u32s: collections.abc.Iterable[builtins.int] | None = ...,
+        u64s: collections.abc.Iterable[builtins.int] | None = ...,
+        f32s: collections.abc.Iterable[builtins.float] | None = ...,
+        f64s: collections.abc.Iterable[builtins.float] | None = ...,
+        c64s: collections.abc.Iterable[builtins.float] | None = ...,
+        c128s: collections.abc.Iterable[builtins.float] | None = ...,
+        tuple_literals: collections.abc.Iterable[global___LiteralProto] | None = ...,
+        f16s: builtins.bytes | None = ...,
+        bf16s: builtins.bytes | None = ...,
+        u16s: builtins.bytes | None = ...,
+        s16s: builtins.bytes | None = ...,
+        f8e5m2s: builtins.bytes | None = ...,
+        f8e4m3fns: builtins.bytes | None = ...,
+        f8e4m3b11fnuzs: builtins.bytes | None = ...,
+        f8e5m2fnuzs: builtins.bytes | None = ...,
+        f8e4m3fnuzs: builtins.bytes | None = ...,
+        sparse_indices: collections.abc.Iterable[builtins.int] | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["shape", b"shape"]) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "bf16s",
+            b"bf16s",
+            "c128s",
+            b"c128s",
+            "c64s",
+            b"c64s",
+            "f16s",
+            b"f16s",
+            "f32s",
+            b"f32s",
+            "f64s",
+            b"f64s",
+            "f8e4m3b11fnuzs",
+            b"f8e4m3b11fnuzs",
+            "f8e4m3fns",
+            b"f8e4m3fns",
+            "f8e4m3fnuzs",
+            b"f8e4m3fnuzs",
+            "f8e5m2fnuzs",
+            b"f8e5m2fnuzs",
+            "f8e5m2s",
+            b"f8e5m2s",
+            "preds",
+            b"preds",
+            "s16s",
+            b"s16s",
+            "s2s",
+            b"s2s",
+            "s32s",
+            b"s32s",
+            "s4s",
+            b"s4s",
+            "s64s",
+            b"s64s",
+            "s8s",
+            b"s8s",
+            "shape",
+            b"shape",
+            "sparse_indices",
+            b"sparse_indices",
+            "tuple_literals",
+            b"tuple_literals",
+            "u16s",
+            b"u16s",
+            "u2s",
+            b"u2s",
+            "u32s",
+            b"u32s",
+            "u4s",
+            b"u4s",
+            "u64s",
+            b"u64s",
+            "u8s",
+            b"u8s",
+        ],
+    ) -> None: ...
+
+global___LiteralProto = LiteralProto
+
+@typing.final
+class WindowDimension(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    SIZE_FIELD_NUMBER: builtins.int
+    STRIDE_FIELD_NUMBER: builtins.int
+    PADDING_LOW_FIELD_NUMBER: builtins.int
+    PADDING_HIGH_FIELD_NUMBER: builtins.int
+    WINDOW_DILATION_FIELD_NUMBER: builtins.int
+    BASE_DILATION_FIELD_NUMBER: builtins.int
+    WINDOW_REVERSAL_FIELD_NUMBER: builtins.int
+    size: builtins.int
+    """The size of the window in this dimension. For a rectangle, this would be
+    the width or height.
+    """
+    stride: builtins.int
+    """The stride at which the window moves across the base area in this
+    dimension. In other words, this is the spacing between different
+    positions of the window in this dimension.
+    """
+    padding_low: builtins.int
+    """If positive, means the amount of padding to add to the base area at the low
+    end of this dimension; if negative, its negative means the number of
+    elements removed from the low end of this dimension. For example, in the
+    horizontal dimension of a rectangle, this would be the number of padding
+    values to pad on the left, given that indices increase when going right.
+    The actual padding value depends upon the context. Convolution pads with
+    zeros. ReduceWindow and SelectAndScatter pads with the reduce function's
+    init value.
+    """
+    padding_high: builtins.int
+    """As padding_low, but on the high end of this dimension. For example, in the
+    horizontal dimension of a rectangle, this would be the number of values to
+    pad on the right, given that indices increase when going right.
+    """
+    window_dilation: builtins.int
+    """Dilation factor of the sliding window in this dimension. A dilation factor
+    of 1 means no dilation. window_dilation - 1 no-op entries ("holes") are
+    implicitly placed between each kernel element. This value may not be less
+    than 1. See documentation for convolution.
+    """
+    base_dilation: builtins.int
+    """Dilation factor of the base area in this dimension. A dilation factor of 1
+    means no dilation. base_dilation - 1 no-op entries ("holes") are implicitly
+    placed between each base area element. This value may not be less than 1.
+    See documentation for convolution.
+    """
+    window_reversal: builtins.bool
+    """Window reversal means that this dimension was logically reversed before the
+    operation.
+    """
+    def __init__(
+        self,
+        *,
+        size: builtins.int | None = ...,
+        stride: builtins.int | None = ...,
+        padding_low: builtins.int | None = ...,
+        padding_high: builtins.int | None = ...,
+        window_dilation: builtins.int | None = ...,
+        base_dilation: builtins.int | None = ...,
+        window_reversal: builtins.bool | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "base_dilation",
+            b"base_dilation",
+            "padding_high",
+            b"padding_high",
+            "padding_low",
+            b"padding_low",
+            "size",
+            b"size",
+            "stride",
+            b"stride",
+            "window_dilation",
+            b"window_dilation",
+            "window_reversal",
+            b"window_reversal",
+        ],
+    ) -> None: ...
+
+global___WindowDimension = WindowDimension
+
+@typing.final
+class Window(google.protobuf.message.Message):
+    """Describes the windowing in an operation such as convolution.
+
+    The window is moved across a base area and for each position of the
+    window a computation is performed. The field below describes the
+    window and the movement of the window across a base area.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    DIMENSIONS_FIELD_NUMBER: builtins.int
+    @property
+    def dimensions(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___WindowDimension]: ...
+    def __init__(self, *, dimensions: collections.abc.Iterable[global___WindowDimension] | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["dimensions", b"dimensions"]) -> None: ...
+
+global___Window = Window
+
+@typing.final
+class GatherDimensionNumbers(google.protobuf.message.Message):
+    """Describes the dimension numbers for a gather operation.
+
+    See https://www.tensorflow.org/performance/xla/operation_semantics#gather for
+    more details.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    OFFSET_DIMS_FIELD_NUMBER: builtins.int
+    COLLAPSED_SLICE_DIMS_FIELD_NUMBER: builtins.int
+    START_INDEX_MAP_FIELD_NUMBER: builtins.int
+    INDEX_VECTOR_DIM_FIELD_NUMBER: builtins.int
+    OPERAND_BATCHING_DIMS_FIELD_NUMBER: builtins.int
+    START_INDICES_BATCHING_DIMS_FIELD_NUMBER: builtins.int
+    index_vector_dim: builtins.int
+    """The dimension in the start_indices input that contains the starting
+    indices.
+    """
+    @property
+    def offset_dims(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """ "Window indices" is a term for a set of indices that index into the
+        interior of a dynamic-slice from the input tensor, the starting indices for
+        which were computed from output_gather_dims (see the operation semantic for
+        how this is defined) and the start_indices tensor.
+
+        The window indices for a specific output index Out is computed as:
+
+         i = 0
+         for (k : [0, input_tensor_shape.rank))
+           window_indices[k] =
+             if k in collapsed_slice_dims
+             then 0
+             else Out[offset_dims[i++]]
+        """
+
+    @property
+    def collapsed_slice_dims(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ...
+    @property
+    def start_index_map(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """This is interpreted as a map from i to start_index_map[i]. It
+        transforms the gather index looked up from the start_indices tensor into
+        the starting index in the input space.
+        """
+
+    @property
+    def operand_batching_dims(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """This is the batch dimensions in the operand."""
+
+    @property
+    def start_indices_batching_dims(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """This is the batch dimensions in the index, and it should be the same size
+        as operand_batching_dims.
+        """
+
+    def __init__(
+        self,
+        *,
+        offset_dims: collections.abc.Iterable[builtins.int] | None = ...,
+        collapsed_slice_dims: collections.abc.Iterable[builtins.int] | None = ...,
+        start_index_map: collections.abc.Iterable[builtins.int] | None = ...,
+        index_vector_dim: builtins.int | None = ...,
+        operand_batching_dims: collections.abc.Iterable[builtins.int] | None = ...,
+        start_indices_batching_dims: collections.abc.Iterable[builtins.int] | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "collapsed_slice_dims",
+            b"collapsed_slice_dims",
+            "index_vector_dim",
+            b"index_vector_dim",
+            "offset_dims",
+            b"offset_dims",
+            "operand_batching_dims",
+            b"operand_batching_dims",
+            "start_index_map",
+            b"start_index_map",
+            "start_indices_batching_dims",
+            b"start_indices_batching_dims",
+        ],
+    ) -> None: ...
+
+global___GatherDimensionNumbers = GatherDimensionNumbers
+
+@typing.final
+class ScatterDimensionNumbers(google.protobuf.message.Message):
+    """Describes the dimension numbers for a scatter operation.
+
+    All the fields are similar to the corresponding fields in
+    GatherDimensionNumbers. Differences are noted below.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    UPDATE_WINDOW_DIMS_FIELD_NUMBER: builtins.int
+    INSERTED_WINDOW_DIMS_FIELD_NUMBER: builtins.int
+    SCATTER_DIMS_TO_OPERAND_DIMS_FIELD_NUMBER: builtins.int
+    INDEX_VECTOR_DIM_FIELD_NUMBER: builtins.int
+    INPUT_BATCHING_DIMS_FIELD_NUMBER: builtins.int
+    SCATTER_INDICES_BATCHING_DIMS_FIELD_NUMBER: builtins.int
+    index_vector_dim: builtins.int
+    @property
+    def update_window_dims(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The set of dimensions in the updates shape that are window dimensions."""
+
+    @property
+    def inserted_window_dims(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The set of window dimensions that must be inserted into the updates shape."""
+
+    @property
+    def scatter_dims_to_operand_dims(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ...
+    @property
+    def input_batching_dims(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """This is the batch dimension in the input."""
+
+    @property
+    def scatter_indices_batching_dims(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """This is the batch dimension in the index."""
+
+    def __init__(
+        self,
+        *,
+        update_window_dims: collections.abc.Iterable[builtins.int] | None = ...,
+        inserted_window_dims: collections.abc.Iterable[builtins.int] | None = ...,
+        scatter_dims_to_operand_dims: collections.abc.Iterable[builtins.int] | None = ...,
+        index_vector_dim: builtins.int | None = ...,
+        input_batching_dims: collections.abc.Iterable[builtins.int] | None = ...,
+        scatter_indices_batching_dims: collections.abc.Iterable[builtins.int] | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "index_vector_dim",
+            b"index_vector_dim",
+            "input_batching_dims",
+            b"input_batching_dims",
+            "inserted_window_dims",
+            b"inserted_window_dims",
+            "scatter_dims_to_operand_dims",
+            b"scatter_dims_to_operand_dims",
+            "scatter_indices_batching_dims",
+            b"scatter_indices_batching_dims",
+            "update_window_dims",
+            b"update_window_dims",
+        ],
+    ) -> None: ...
+
+global___ScatterDimensionNumbers = ScatterDimensionNumbers
+
+@typing.final
+class ConvolutionDimensionNumbers(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    INPUT_BATCH_DIMENSION_FIELD_NUMBER: builtins.int
+    INPUT_FEATURE_DIMENSION_FIELD_NUMBER: builtins.int
+    INPUT_SPATIAL_DIMENSIONS_FIELD_NUMBER: builtins.int
+    KERNEL_INPUT_FEATURE_DIMENSION_FIELD_NUMBER: builtins.int
+    KERNEL_OUTPUT_FEATURE_DIMENSION_FIELD_NUMBER: builtins.int
+    KERNEL_SPATIAL_DIMENSIONS_FIELD_NUMBER: builtins.int
+    OUTPUT_BATCH_DIMENSION_FIELD_NUMBER: builtins.int
+    OUTPUT_FEATURE_DIMENSION_FIELD_NUMBER: builtins.int
+    OUTPUT_SPATIAL_DIMENSIONS_FIELD_NUMBER: builtins.int
+    input_batch_dimension: builtins.int
+    """The number of the dimension that represents batch in the input."""
+    input_feature_dimension: builtins.int
+    """The number of the dimension that represents features in the input."""
+    kernel_input_feature_dimension: builtins.int
+    """The number of the dimension that represents input features in the
+    convolutional kernel (rhs).
+    """
+    kernel_output_feature_dimension: builtins.int
+    """The number of the dimension that represents output features in
+    the convolutional kernel (rhs).
+    """
+    output_batch_dimension: builtins.int
+    """The number of the dimension that represents batch in the output."""
+    output_feature_dimension: builtins.int
+    """The number of the dimension that represents features in the output."""
+    @property
+    def input_spatial_dimensions(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The dimension numbers for the spatial dimensions that the window
+        moves through in the input.
+        """
+
+    @property
+    def kernel_spatial_dimensions(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The dimension numbers for the spatial dimensions that the window
+        moves through in the kernel (rhs). window.strides(0) is the
+        stride in the kernel_spatial_dimensions(0) dimension.
+        """
+
+    @property
+    def output_spatial_dimensions(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The dimension numbers for the spatial dimensions that the window
+        moves through in the output.
+        """
+
+    def __init__(
+        self,
+        *,
+        input_batch_dimension: builtins.int | None = ...,
+        input_feature_dimension: builtins.int | None = ...,
+        input_spatial_dimensions: collections.abc.Iterable[builtins.int] | None = ...,
+        kernel_input_feature_dimension: builtins.int | None = ...,
+        kernel_output_feature_dimension: builtins.int | None = ...,
+        kernel_spatial_dimensions: collections.abc.Iterable[builtins.int] | None = ...,
+        output_batch_dimension: builtins.int | None = ...,
+        output_feature_dimension: builtins.int | None = ...,
+        output_spatial_dimensions: collections.abc.Iterable[builtins.int] | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "input_batch_dimension",
+            b"input_batch_dimension",
+            "input_feature_dimension",
+            b"input_feature_dimension",
+            "input_spatial_dimensions",
+            b"input_spatial_dimensions",
+            "kernel_input_feature_dimension",
+            b"kernel_input_feature_dimension",
+            "kernel_output_feature_dimension",
+            b"kernel_output_feature_dimension",
+            "kernel_spatial_dimensions",
+            b"kernel_spatial_dimensions",
+            "output_batch_dimension",
+            b"output_batch_dimension",
+            "output_feature_dimension",
+            b"output_feature_dimension",
+            "output_spatial_dimensions",
+            b"output_spatial_dimensions",
+        ],
+    ) -> None: ...
+
+global___ConvolutionDimensionNumbers = ConvolutionDimensionNumbers
+
+@typing.final
+class DotDimensionNumbers(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    LHS_CONTRACTING_DIMENSIONS_FIELD_NUMBER: builtins.int
+    RHS_CONTRACTING_DIMENSIONS_FIELD_NUMBER: builtins.int
+    LHS_BATCH_DIMENSIONS_FIELD_NUMBER: builtins.int
+    RHS_BATCH_DIMENSIONS_FIELD_NUMBER: builtins.int
+    @property
+    def lhs_contracting_dimensions(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The dimension numbers that represent the 'lhs' contracting dimensions."""
+
+    @property
+    def rhs_contracting_dimensions(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The dimension numbers that represent the 'rhs' contracting dimensions."""
+
+    @property
+    def lhs_batch_dimensions(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The dimension numbers that represent the 'lhs' batch dimensions."""
+
+    @property
+    def rhs_batch_dimensions(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The dimension numbers that represent the 'rhs' batch dimensions."""
+
+    def __init__(
+        self,
+        *,
+        lhs_contracting_dimensions: collections.abc.Iterable[builtins.int] | None = ...,
+        rhs_contracting_dimensions: collections.abc.Iterable[builtins.int] | None = ...,
+        lhs_batch_dimensions: collections.abc.Iterable[builtins.int] | None = ...,
+        rhs_batch_dimensions: collections.abc.Iterable[builtins.int] | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "lhs_batch_dimensions",
+            b"lhs_batch_dimensions",
+            "lhs_contracting_dimensions",
+            b"lhs_contracting_dimensions",
+            "rhs_batch_dimensions",
+            b"rhs_batch_dimensions",
+            "rhs_contracting_dimensions",
+            b"rhs_contracting_dimensions",
+        ],
+    ) -> None: ...
+
+global___DotDimensionNumbers = DotDimensionNumbers
+
+@typing.final
+class SparsityDescriptor(google.protobuf.message.Message):
+    """Contains sparsity metadata for a sparse dot operation.
+    The only supported type atm is structured 2:4 sparsity, which is natively
+    supported on NVidia GPUs.
+    Restrictions:
+    - only one operand of the dot operation may be sparse;
+    - only the contracting dimension may be sparse.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    TYPE_FIELD_NUMBER: builtins.int
+    INDEX_FIELD_NUMBER: builtins.int
+    DIMENSION_FIELD_NUMBER: builtins.int
+    N_FIELD_NUMBER: builtins.int
+    M_FIELD_NUMBER: builtins.int
+    type: global___SparsityType.ValueType
+    index: builtins.int
+    """Sparse operand index (0 or 1)."""
+    dimension: builtins.int
+    """Sparse dimension number."""
+    n: builtins.int
+    """Structured N:M sparsity (N < M)."""
+    m: builtins.int
+    def __init__(
+        self,
+        *,
+        type: global___SparsityType.ValueType | None = ...,
+        index: builtins.int | None = ...,
+        dimension: builtins.int | None = ...,
+        n: builtins.int | None = ...,
+        m: builtins.int | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self, field_name: typing.Literal["dimension", b"dimension", "index", b"index", "m", b"m", "n", b"n", "type", b"type"]
+    ) -> None: ...
+
+global___SparsityDescriptor = SparsityDescriptor
+
+@typing.final
+class TriangularSolveOptions(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    class _Transpose:
+        ValueType = typing.NewType("ValueType", builtins.int)
+        V: typing_extensions.TypeAlias = ValueType
+
+    class _TransposeEnumTypeWrapper(
+        google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[TriangularSolveOptions._Transpose.ValueType], builtins.type
+    ):
+        DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+        TRANSPOSE_INVALID: TriangularSolveOptions._Transpose.ValueType  # 0
+        NO_TRANSPOSE: TriangularSolveOptions._Transpose.ValueType  # 1
+        """Don't transpose 'a'."""
+        TRANSPOSE: TriangularSolveOptions._Transpose.ValueType  # 2
+        """Transpose 'a'."""
+        ADJOINT: TriangularSolveOptions._Transpose.ValueType  # 3
+        """Complex conjugate and transpose 'a'."""
+
+    class Transpose(_Transpose, metaclass=_TransposeEnumTypeWrapper):
+        """Should we transpose or use the adjoint of 'a'?"""
+
+    TRANSPOSE_INVALID: TriangularSolveOptions.Transpose.ValueType  # 0
+    NO_TRANSPOSE: TriangularSolveOptions.Transpose.ValueType  # 1
+    """Don't transpose 'a'."""
+    TRANSPOSE: TriangularSolveOptions.Transpose.ValueType  # 2
+    """Transpose 'a'."""
+    ADJOINT: TriangularSolveOptions.Transpose.ValueType  # 3
+    """Complex conjugate and transpose 'a'."""
+
+    LEFT_SIDE_FIELD_NUMBER: builtins.int
+    LOWER_FIELD_NUMBER: builtins.int
+    UNIT_DIAGONAL_FIELD_NUMBER: builtins.int
+    TRANSPOSE_A_FIELD_NUMBER: builtins.int
+    left_side: builtins.bool
+    """If true, solves ax = b. If false, solves xa = b."""
+    lower: builtins.bool
+    """If true, 'a' is lower triangular. If false, 'a' is upper triangular."""
+    unit_diagonal: builtins.bool
+    """If true, the diagonal elements of 'a' are assumed to be 1 and not accessed."""
+    transpose_a: global___TriangularSolveOptions.Transpose.ValueType
+    def __init__(
+        self,
+        *,
+        left_side: builtins.bool | None = ...,
+        lower: builtins.bool | None = ...,
+        unit_diagonal: builtins.bool | None = ...,
+        transpose_a: global___TriangularSolveOptions.Transpose.ValueType | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "left_side", b"left_side", "lower", b"lower", "transpose_a", b"transpose_a", "unit_diagonal", b"unit_diagonal"
+        ],
+    ) -> None: ...
+
+global___TriangularSolveOptions = TriangularSolveOptions
+
+@typing.final
+class CholeskyOptions(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    LOWER_FIELD_NUMBER: builtins.int
+    lower: builtins.bool
+    """If true, uses the lower triangle of `a`. If false, uses the upper triangle
+    of `a`.
+    """
+    def __init__(self, *, lower: builtins.bool | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["lower", b"lower"]) -> None: ...
+
+global___CholeskyOptions = CholeskyOptions
+
+@typing.final
+class SortOptions(google.protobuf.message.Message):
+    """Attributes of the sort custom call (cub::DeviceRadixSort)."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    DESCENDING_FIELD_NUMBER: builtins.int
+    descending: builtins.bool
+    def __init__(self, *, descending: builtins.bool | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["descending", b"descending"]) -> None: ...
+
+global___SortOptions = SortOptions
+
+@typing.final
+class FrontendAttributes(google.protobuf.message.Message):
+    """Generic map of attributes used to pass hints / configuration options from
+    the Python frontend to the XLA backend.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.final
+    class MapEntry(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        KEY_FIELD_NUMBER: builtins.int
+        VALUE_FIELD_NUMBER: builtins.int
+        key: builtins.str
+        value: builtins.str
+        def __init__(self, *, key: builtins.str | None = ..., value: builtins.str | None = ...) -> None: ...
+        def ClearField(self, field_name: typing.Literal["key", b"key", "value", b"value"]) -> None: ...
+
+    MAP_FIELD_NUMBER: builtins.int
+    @property
+    def map(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: ...
+    def __init__(self, *, map: collections.abc.Mapping[builtins.str, builtins.str] | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["map", b"map"]) -> None: ...
+
+global___FrontendAttributes = FrontendAttributes
+
+@typing.final
+class Statistic(google.protobuf.message.Message):
+    """Represents a single statistic to track."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    STAT_NAME_FIELD_NUMBER: builtins.int
+    STAT_VAL_FIELD_NUMBER: builtins.int
+    stat_name: builtins.str
+    """Must be a single word consisting of any alphanumeric characters"""
+    stat_val: builtins.float
+    """Must be within a range of [0, 100], in order for the graph dumper to
+    properly render the statistic onto the graph.
+    """
+    def __init__(self, *, stat_name: builtins.str | None = ..., stat_val: builtins.float | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["stat_name", b"stat_name", "stat_val", b"stat_val"]) -> None: ...
+
+global___Statistic = Statistic
+
+@typing.final
+class StatisticsViz(google.protobuf.message.Message):
+    """Represents the information needed to visualize propagation statistics when
+    rendering an HLO graph. This includes an array of statistics as well as the
+    index of the statistic to render.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    STAT_INDEX_TO_VISUALIZE_FIELD_NUMBER: builtins.int
+    STATISTICS_FIELD_NUMBER: builtins.int
+    stat_index_to_visualize: builtins.int
+    @property
+    def statistics(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Statistic]: ...
+    def __init__(
+        self,
+        *,
+        stat_index_to_visualize: builtins.int | None = ...,
+        statistics: collections.abc.Iterable[global___Statistic] | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self, field_name: typing.Literal["stat_index_to_visualize", b"stat_index_to_visualize", "statistics", b"statistics"]
+    ) -> None: ...
+
+global___StatisticsViz = StatisticsViz
+
+@typing.final
+class OpSharding(google.protobuf.message.Message):
+    """LINT.IfChange"""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    class _Type:
+        ValueType = typing.NewType("ValueType", builtins.int)
+        V: typing_extensions.TypeAlias = ValueType
+
+    class _TypeEnumTypeWrapper(
+        google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[OpSharding._Type.ValueType], builtins.type
+    ):
+        DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+        REPLICATED: OpSharding._Type.ValueType  # 0
+        """This sharding is replicated across all devices (implies maximal,
+        all other fields are unused).
+        """
+        MAXIMAL: OpSharding._Type.ValueType  # 1
+        """This sharding is maximal - one device runs the entire operation."""
+        TUPLE: OpSharding._Type.ValueType  # 2
+        """This sharding is a tuple - only the tuple_shardings field is valid."""
+        OTHER: OpSharding._Type.ValueType  # 3
+        """None of the above; tile_shape and tile_assignment are both used."""
+        MANUAL: OpSharding._Type.ValueType  # 4
+        """This op is manually sharded: the shapes are already partitioned and the
+        partitioner should not change this op.
+        """
+        UNKNOWN: OpSharding._Type.ValueType  # 5
+        """This sharding is a placeholder sharding with lowest precedence, it can be
+        overwriten by any other shardings.
+        """
+
+    class Type(_Type, metaclass=_TypeEnumTypeWrapper): ...
+    REPLICATED: OpSharding.Type.ValueType  # 0
+    """This sharding is replicated across all devices (implies maximal,
+    all other fields are unused).
+    """
+    MAXIMAL: OpSharding.Type.ValueType  # 1
+    """This sharding is maximal - one device runs the entire operation."""
+    TUPLE: OpSharding.Type.ValueType  # 2
+    """This sharding is a tuple - only the tuple_shardings field is valid."""
+    OTHER: OpSharding.Type.ValueType  # 3
+    """None of the above; tile_shape and tile_assignment are both used."""
+    MANUAL: OpSharding.Type.ValueType  # 4
+    """This op is manually sharded: the shapes are already partitioned and the
+    partitioner should not change this op.
+    """
+    UNKNOWN: OpSharding.Type.ValueType  # 5
+    """This sharding is a placeholder sharding with lowest precedence, it can be
+    overwriten by any other shardings.
+    """
+
+    class _ShardGroupType:
+        ValueType = typing.NewType("ValueType", builtins.int)
+        V: typing_extensions.TypeAlias = ValueType
+
+    class _ShardGroupTypeEnumTypeWrapper(
+        google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[OpSharding._ShardGroupType.ValueType], builtins.type
+    ):
+        DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+        AS: OpSharding._ShardGroupType.ValueType  # 0
+        """This op will be sharded exactly the same as the other op. (hard
+        restriction)
+        """
+        LIKE: OpSharding._ShardGroupType.ValueType  # 1
+        """This op will try to allow sharding propagation within the same group even
+        there is no data dependencies among them, but there is no guarantee that
+        the final shardings within the same group will be exactly the same. (soft
+        restriction)
+        """
+
+    class ShardGroupType(_ShardGroupType, metaclass=_ShardGroupTypeEnumTypeWrapper):
+        """Used to decide whether this op is to be sharded like some other ops, or to
+        which other ops will be sharded like.
+        """
+
+    AS: OpSharding.ShardGroupType.ValueType  # 0
+    """This op will be sharded exactly the same as the other op. (hard
+    restriction)
+    """
+    LIKE: OpSharding.ShardGroupType.ValueType  # 1
+    """This op will try to allow sharding propagation within the same group even
+    there is no data dependencies among them, but there is no guarantee that
+    the final shardings within the same group will be exactly the same. (soft
+    restriction)
+    """
+
+    TYPE_FIELD_NUMBER: builtins.int
+    TILE_SHAPE_FIELD_NUMBER: builtins.int
+    TILE_ASSIGNMENT_DIMENSIONS_FIELD_NUMBER: builtins.int
+    TILE_ASSIGNMENT_DEVICES_FIELD_NUMBER: builtins.int
+    TUPLE_SHARDINGS_FIELD_NUMBER: builtins.int
+    REPLICATE_ON_LAST_TILE_DIM_FIELD_NUMBER: builtins.int
+    METADATA_FIELD_NUMBER: builtins.int
+    LAST_TILE_DIMS_FIELD_NUMBER: builtins.int
+    IOTA_RESHAPE_DIMS_FIELD_NUMBER: builtins.int
+    IOTA_TRANSPOSE_PERM_FIELD_NUMBER: builtins.int
+    IS_SHARD_GROUP_FIELD_NUMBER: builtins.int
+    SHARD_GROUP_ID_FIELD_NUMBER: builtins.int
+    SHARD_GROUP_TYPE_FIELD_NUMBER: builtins.int
+    type: global___OpSharding.Type.ValueType
+    replicate_on_last_tile_dim: builtins.bool
+    """Only used for OTHER type. If true, data is sharded according to other
+    dimensions of tile_assignment(), but replicated across devices along the
+    last dimension. (Experimental)
+    """
+    is_shard_group: builtins.bool
+    """This field decides whether this op is in a shard group."""
+    shard_group_id: builtins.int
+    """This field is used to store the unique id of the shard group."""
+    shard_group_type: global___OpSharding.ShardGroupType.ValueType
+    @property
+    def tile_shape(self) -> global___ShapeProto:
+        """The shape of the sharded tile."""
+
+    @property
+    def tile_assignment_dimensions(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The shape of the tile assignment tensor - this must be the same rank as
+        tile_shape and the product of its dimensions must equal
+        tile_assignment_devices.size().
+        """
+
+    @property
+    def tile_assignment_devices(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """Flattened list of device IDs. The order of flattening is the same as used
+        by IndexUtil::MultiToLinearIndex(tile_assignment_shape).
+        Only one of tile_assignment_devices and iota_dimensions shall be non-empty.
+        """
+
+    @property
+    def tuple_shardings(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___OpSharding]:
+        """If type == TUPLE, the sub-shardings, one per leaf node in the tuple shape,
+        in pre-order. The tuple shape could be nested; here we store just a
+        flattened list of all leaves in the tuple shape. Note that the tuple shape
+        is not stored here; shardings do not store the shapes to which they are
+        applied, this is inferred from the instruction this sharding gets attached
+        to.
+        """
+
+    @property
+    def metadata(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___OpMetadata]:
+        """This field is used to track the source of this sharding, usually derived
+        from instructions. Multple metadata may be populated if sharding is
+        combined with other shardings.  Metadata are to not be populated when
+        type == TUPLE and instead metadata should be set on individual tuple
+        elements.
+        """
+
+    @property
+    def last_tile_dims(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[global___OpSharding.Type.ValueType]:
+        """This field is used to represented the sharding type of each subgroup.
+        For example, sharding={devices=[2,2,2,2]0,1,2,...,15 last_tile_dims={
+        replicate, manual, unreduced}} means that each of the last 3 dimensions
+        in [2,2,2,2] represents a subgrouping in replicate, manual,
+        unreduced sharding type respectively.
+        """
+
+    @property
+    def iota_reshape_dims(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """Dimensions used to reshape the 1D iota array of device IDs.
+        Only one of tile_assignment_devices and iota_reshape_dims shall be
+        non-empty.
+        """
+
+    @property
+    def iota_transpose_perm(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """Dimension permutations to transposed the iota array reshaped to
+        iota_reshape_dims. This must have the same size as iota_reshape_dims.
+        """
+
+    def __init__(
+        self,
+        *,
+        type: global___OpSharding.Type.ValueType | None = ...,
+        tile_shape: global___ShapeProto | None = ...,
+        tile_assignment_dimensions: collections.abc.Iterable[builtins.int] | None = ...,
+        tile_assignment_devices: collections.abc.Iterable[builtins.int] | None = ...,
+        tuple_shardings: collections.abc.Iterable[global___OpSharding] | None = ...,
+        replicate_on_last_tile_dim: builtins.bool | None = ...,
+        metadata: collections.abc.Iterable[global___OpMetadata] | None = ...,
+        last_tile_dims: collections.abc.Iterable[global___OpSharding.Type.ValueType] | None = ...,
+        iota_reshape_dims: collections.abc.Iterable[builtins.int] | None = ...,
+        iota_transpose_perm: collections.abc.Iterable[builtins.int] | None = ...,
+        is_shard_group: builtins.bool | None = ...,
+        shard_group_id: builtins.int | None = ...,
+        shard_group_type: global___OpSharding.ShardGroupType.ValueType | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["tile_shape", b"tile_shape"]) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "iota_reshape_dims",
+            b"iota_reshape_dims",
+            "iota_transpose_perm",
+            b"iota_transpose_perm",
+            "is_shard_group",
+            b"is_shard_group",
+            "last_tile_dims",
+            b"last_tile_dims",
+            "metadata",
+            b"metadata",
+            "replicate_on_last_tile_dim",
+            b"replicate_on_last_tile_dim",
+            "shard_group_id",
+            b"shard_group_id",
+            "shard_group_type",
+            b"shard_group_type",
+            "tile_assignment_devices",
+            b"tile_assignment_devices",
+            "tile_assignment_dimensions",
+            b"tile_assignment_dimensions",
+            "tile_shape",
+            b"tile_shape",
+            "tuple_shardings",
+            b"tuple_shardings",
+            "type",
+            b"type",
+        ],
+    ) -> None: ...
+
+global___OpSharding = OpSharding
+
+@typing.final
+class ReplicaGroup(google.protobuf.message.Message):
+    """Describes the replica groups in a cross replica op (e.g., all-reduce and
+    all-to-all).
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    REPLICA_IDS_FIELD_NUMBER: builtins.int
+    @property
+    def replica_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The ids of the replicas that belongs to the same group. The ordering of the
+        ids matters in some ops (e.g., all-to-all).
+        """
+
+    def __init__(self, *, replica_ids: collections.abc.Iterable[builtins.int] | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["replica_ids", b"replica_ids"]) -> None: ...
+
+global___ReplicaGroup = ReplicaGroup
+
+@typing.final
+class IotaReplicaGroupListProto(google.protobuf.message.Message):
+    """Represents a list of replica groups (a list of list of devices) with
+    reshaping and transposing an iota array (iota tile assignment). Can be used
+    to represent certain common patterns of device lists in a compact, scalable
+    format.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    NUM_REPLICA_GROUPS_FIELD_NUMBER: builtins.int
+    NUM_DEVICES_PER_GROUP_FIELD_NUMBER: builtins.int
+    IOTA_RESHAPE_DIMS_FIELD_NUMBER: builtins.int
+    IOTA_TRANSPOSE_PERM_FIELD_NUMBER: builtins.int
+    num_replica_groups: builtins.int
+    """Number of replica groups."""
+    num_devices_per_group: builtins.int
+    """Number of devices per group."""
+    @property
+    def iota_reshape_dims(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The dimensions used to reshape the 1D iota array of device IDs."""
+
+    @property
+    def iota_transpose_perm(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
+        """The dimension permutations to transposed the iota array reshaped to
+        iota_reshape_dims. This must have the same size as iota_reshape_dims.
+        """
+
+    def __init__(
+        self,
+        *,
+        num_replica_groups: builtins.int | None = ...,
+        num_devices_per_group: builtins.int | None = ...,
+        iota_reshape_dims: collections.abc.Iterable[builtins.int] | None = ...,
+        iota_transpose_perm: collections.abc.Iterable[builtins.int] | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "iota_reshape_dims",
+            b"iota_reshape_dims",
+            "iota_transpose_perm",
+            b"iota_transpose_perm",
+            "num_devices_per_group",
+            b"num_devices_per_group",
+            "num_replica_groups",
+            b"num_replica_groups",
+        ],
+    ) -> None: ...
+
+global___IotaReplicaGroupListProto = IotaReplicaGroupListProto
+
+@typing.final
+class CollectiveDeviceListProto(google.protobuf.message.Message):
+    """Represents a series of devices participating in a collective operation (e.g.,
+    all-reduce and all-to-all). While this directly translates to a list of
+    replica groups, it may be used to represent these lists in a compact form.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    REPLICA_GROUPS_FIELD_NUMBER: builtins.int
+    IOTA_REPLICA_GROUP_LIST_FIELD_NUMBER: builtins.int
+    @property
+    def replica_groups(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ReplicaGroup]:
+        """ReplicaGroupV1: List of replica groups. Legacy way of representing device
+        lists.
+        """
+
+    @property
+    def iota_replica_group_list(self) -> global___IotaReplicaGroupListProto:
+        """ReplicaGroupV2: Represents a list of replica groups with reshaping and
+        transposing an iota array.
+        """
+
+    def __init__(
+        self,
+        *,
+        replica_groups: collections.abc.Iterable[global___ReplicaGroup] | None = ...,
+        iota_replica_group_list: global___IotaReplicaGroupListProto | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["iota_replica_group_list", b"iota_replica_group_list"]) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal["iota_replica_group_list", b"iota_replica_group_list", "replica_groups", b"replica_groups"],
+    ) -> None: ...
+
+global___CollectiveDeviceListProto = CollectiveDeviceListProto
+
+@typing.final
+class SourceTarget(google.protobuf.message.Message):
+    """Describes the source target pair in the collective permute op."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    SOURCE_FIELD_NUMBER: builtins.int
+    TARGET_FIELD_NUMBER: builtins.int
+    source: builtins.int
+    target: builtins.int
+    def __init__(self, *, source: builtins.int | None = ..., target: builtins.int | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["source", b"source", "target", b"target"]) -> None: ...
+
+global___SourceTarget = SourceTarget
+
+@typing.final
+class PrecisionConfig(google.protobuf.message.Message):
+    """Used to indicate the precision configuration. It has backend specific
+    meaning.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    class _Precision:
+        ValueType = typing.NewType("ValueType", builtins.int)
+        V: typing_extensions.TypeAlias = ValueType
+
+    class _PrecisionEnumTypeWrapper(
+        google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[PrecisionConfig._Precision.ValueType], builtins.type
+    ):
+        DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+        DEFAULT: PrecisionConfig._Precision.ValueType  # 0
+        HIGH: PrecisionConfig._Precision.ValueType  # 1
+        HIGHEST: PrecisionConfig._Precision.ValueType  # 2
+        PACKED_NIBBLE: PrecisionConfig._Precision.ValueType  # 3
+        """Each U8/S8 value in a tensor actually represents 2 nibble values."""
+
+    class Precision(_Precision, metaclass=_PrecisionEnumTypeWrapper): ...
+    DEFAULT: PrecisionConfig.Precision.ValueType  # 0
+    HIGH: PrecisionConfig.Precision.ValueType  # 1
+    HIGHEST: PrecisionConfig.Precision.ValueType  # 2
+    PACKED_NIBBLE: PrecisionConfig.Precision.ValueType  # 3
+    """Each U8/S8 value in a tensor actually represents 2 nibble values."""
+
+    class _Algorithm:
+        ValueType = typing.NewType("ValueType", builtins.int)
+        V: typing_extensions.TypeAlias = ValueType
+
+    class _AlgorithmEnumTypeWrapper(
+        google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[PrecisionConfig._Algorithm.ValueType], builtins.type
+    ):
+        DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+        ALG_UNSET: PrecisionConfig._Algorithm.ValueType  # 0
+        """If the algorithm is `ALG_UNSET`, we will decide the algorithm based on
+        the operand_precision values (for now).
+        """
+        ALG_DOT_ANY_F8_ANY_F8_F32: PrecisionConfig._Algorithm.ValueType  # 1
+        """The storage type can be any 8-bit floating point type."""
+        ALG_DOT_ANY_F8_ANY_F8_F32_FAST_ACCUM: PrecisionConfig._Algorithm.ValueType  # 2
+        """The storage type can be any 8-bit floating point type. Intermediate
+        results will not periodically be promoted to a higher precision. This
+        corresponds to CUBLASLT_MATMUL_DESC_FAST_ACCUM. Triton's
+        maxNumImpreciseAcc=32 setting may be similar.
+        """
+        ALG_DOT_F16_F16_F16: PrecisionConfig._Algorithm.ValueType  # 3
+        ALG_DOT_F16_F16_F32: PrecisionConfig._Algorithm.ValueType  # 4
+        ALG_DOT_BF16_BF16_BF16: PrecisionConfig._Algorithm.ValueType  # 5
+        ALG_DOT_BF16_BF16_F32: PrecisionConfig._Algorithm.ValueType  # 6
+        ALG_DOT_BF16_BF16_F32_X3: PrecisionConfig._Algorithm.ValueType  # 7
+        """An algorithm which uses 3 BF16_BF16_F32 matmuls to achieve better
+        precision.
+        """
+        ALG_DOT_BF16_BF16_F32_X6: PrecisionConfig._Algorithm.ValueType  # 8
+        """An algorithm which uses 6 BF16_BF16_F32 matmuls to achieve better
+        precision (similar to F32).
+        """
+        ALG_DOT_TF32_TF32_F32: PrecisionConfig._Algorithm.ValueType  # 9
+        ALG_DOT_TF32_TF32_F32_X3: PrecisionConfig._Algorithm.ValueType  # 10
+        """An algorithm which uses 3 TF32_TF32_F32 matmuls to achieve better
+        precision (similar to F32).
+        """
+        ALG_DOT_F32_F32_F32: PrecisionConfig._Algorithm.ValueType  # 11
+        ALG_DOT_F64_F64_F64: PrecisionConfig._Algorithm.ValueType  # 12
+
+    class Algorithm(_Algorithm, metaclass=_AlgorithmEnumTypeWrapper):
+        """The algorithm used to evaluate the instruction.
+
+        The naming convention for the dot instruction is
+        ALG_DOT_{A_TYPE}_{B_TYPE}_{ACCUM_TYPE}[_X{NUM_OPS}] where A_TYPE, B_TYPE
+        and ACCUM_TYPE correspond to the types in the "primitive dot operations"
+        (such as TensorCore operations) and NUM_OPS is the number of such
+        operations used per "primitive tile". When the NUM_OPS
+        field is skipped, it is assumed to be 1. The types mentioned in the name
+        are independent of the storage types.
+
+        In general ATYPE and BTYPE are the precisions that the LHS and RHS of the
+        operation are rounded to and ACCUMTYPE is the accumulation type. If a
+        backend does not support the given algorithm, an error is raised. The
+        Algorithm enum is intended to eventually replace the Precision enum.
+        """
+
+    ALG_UNSET: PrecisionConfig.Algorithm.ValueType  # 0
+    """If the algorithm is `ALG_UNSET`, we will decide the algorithm based on
+    the operand_precision values (for now).
+    """
+    ALG_DOT_ANY_F8_ANY_F8_F32: PrecisionConfig.Algorithm.ValueType  # 1
+    """The storage type can be any 8-bit floating point type."""
+    ALG_DOT_ANY_F8_ANY_F8_F32_FAST_ACCUM: PrecisionConfig.Algorithm.ValueType  # 2
+    """The storage type can be any 8-bit floating point type. Intermediate
+    results will not periodically be promoted to a higher precision. This
+    corresponds to CUBLASLT_MATMUL_DESC_FAST_ACCUM. Triton's
+    maxNumImpreciseAcc=32 setting may be similar.
+    """
+    ALG_DOT_F16_F16_F16: PrecisionConfig.Algorithm.ValueType  # 3
+    ALG_DOT_F16_F16_F32: PrecisionConfig.Algorithm.ValueType  # 4
+    ALG_DOT_BF16_BF16_BF16: PrecisionConfig.Algorithm.ValueType  # 5
+    ALG_DOT_BF16_BF16_F32: PrecisionConfig.Algorithm.ValueType  # 6
+    ALG_DOT_BF16_BF16_F32_X3: PrecisionConfig.Algorithm.ValueType  # 7
+    """An algorithm which uses 3 BF16_BF16_F32 matmuls to achieve better
+    precision.
+    """
+    ALG_DOT_BF16_BF16_F32_X6: PrecisionConfig.Algorithm.ValueType  # 8
+    """An algorithm which uses 6 BF16_BF16_F32 matmuls to achieve better
+    precision (similar to F32).
+    """
+    ALG_DOT_TF32_TF32_F32: PrecisionConfig.Algorithm.ValueType  # 9
+    ALG_DOT_TF32_TF32_F32_X3: PrecisionConfig.Algorithm.ValueType  # 10
+    """An algorithm which uses 3 TF32_TF32_F32 matmuls to achieve better
+    precision (similar to F32).
+    """
+    ALG_DOT_F32_F32_F32: PrecisionConfig.Algorithm.ValueType  # 11
+    ALG_DOT_F64_F64_F64: PrecisionConfig.Algorithm.ValueType  # 12
+
+    OPERAND_PRECISION_FIELD_NUMBER: builtins.int
+    ALGORITHM_FIELD_NUMBER: builtins.int
+    algorithm: global___PrecisionConfig.Algorithm.ValueType
+    """Currently doesn't do anything, but we plan to support it for dot and
+    possibly more instructions.
+
+    TODO(b/316147294): Support this on GPU and add this to StableHLO as well.
+
+    If this is set, then `operand_precision` should be set to DEFAULT and it
+    will be ignored.
+    """
+    @property
+    def operand_precision(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[global___PrecisionConfig.Precision.ValueType]: ...
+    def __init__(
+        self,
+        *,
+        operand_precision: collections.abc.Iterable[global___PrecisionConfig.Precision.ValueType] | None = ...,
+        algorithm: global___PrecisionConfig.Algorithm.ValueType | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self, field_name: typing.Literal["algorithm", b"algorithm", "operand_precision", b"operand_precision"]
+    ) -> None: ...
+
+global___PrecisionConfig = PrecisionConfig
+
+@typing.final
+class ParameterReplication(google.protobuf.message.Message):
+    """Describes whether all data-parallelism replicas will receive the same
+    parameter data at each buffer.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    REPLICATED_AT_LEAF_BUFFERS_FIELD_NUMBER: builtins.int
+    @property
+    def replicated_at_leaf_buffers(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.bool]:
+        """A list of boolean values for the flattened leaf buffers. Each value
+        indicates whether the corresponding leaf buffer is replicated.
+
+        If this field is empty, it means no buffer is replicated. Otherwise, the
+        number of elements in this field must match the number of leaf buffers in
+        the HLO instruction's shape.
+        """
+
+    def __init__(self, *, replicated_at_leaf_buffers: collections.abc.Iterable[builtins.bool] | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["replicated_at_leaf_buffers", b"replicated_at_leaf_buffers"]) -> None: ...
+
+global___ParameterReplication = ParameterReplication
+
+@typing.final
+class WhileLoopBackendConfig(google.protobuf.message.Message):
+    """A backend-config for kWhile loops that stores the loop's trip count, if it is
+    known.
+
+    This is useful for backends that can implement a `for i in 0..N` loop more
+    efficiently than a `while` loop.  For example, on GPUs, we can implement a
+    `for i in 0..N` loop by enqueueing the kernels for the loop body N times,
+    whereas implementing a `while` loop requires a host-device sync on each
+    iteration.
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    @typing.final
+    class KnownTripCount(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        N_FIELD_NUMBER: builtins.int
+        n: builtins.int
+        def __init__(self, *, n: builtins.int | None = ...) -> None: ...
+        def ClearField(self, field_name: typing.Literal["n", b"n"]) -> None: ...
+
+    KNOWN_TRIP_COUNT_FIELD_NUMBER: builtins.int
+    @property
+    def known_trip_count(self) -> global___WhileLoopBackendConfig.KnownTripCount:
+        """This indirection lets us distinguish between known-trip-count == 0 and
+        unknown-trip-count.
+        """
+
+    def __init__(self, *, known_trip_count: global___WhileLoopBackendConfig.KnownTripCount | None = ...) -> None: ...
+    def HasField(self, field_name: typing.Literal["known_trip_count", b"known_trip_count"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing.Literal["known_trip_count", b"known_trip_count"]) -> None: ...
+
+global___WhileLoopBackendConfig = WhileLoopBackendConfig
+
+@typing.final
+class OutputOperandAliasing(google.protobuf.message.Message):
+    """Specifies a pair of output/operand buffers that alias each other for
+    kCustomCall and kFusion
+    """
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    OUTPUT_SHAPE_INDEX_FIELD_NUMBER: builtins.int
+    OPERAND_INDEX_FIELD_NUMBER: builtins.int
+    OPERAND_SHAPE_INDEX_FIELD_NUMBER: builtins.int
+    operand_index: builtins.int
+    @property
+    def output_shape_index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ...
+    @property
+    def operand_shape_index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ...
+    def __init__(
+        self,
+        *,
+        output_shape_index: collections.abc.Iterable[builtins.int] | None = ...,
+        operand_index: builtins.int | None = ...,
+        operand_shape_index: collections.abc.Iterable[builtins.int] | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "operand_index",
+            b"operand_index",
+            "operand_shape_index",
+            b"operand_shape_index",
+            "output_shape_index",
+            b"output_shape_index",
+        ],
+    ) -> None: ...
+
+global___OutputOperandAliasing = OutputOperandAliasing
+
+@typing.final
+class OriginalArrayProto(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    LEAF_SHAPE_INDEX_FIELD_NUMBER: builtins.int
+    INSTRUCTION_NAME_FIELD_NUMBER: builtins.int
+    SHAPE_INDEX_FIELD_NUMBER: builtins.int
+    instruction_name: builtins.str
+    @property
+    def leaf_shape_index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ...
+    @property
+    def shape_index(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ...
+    def __init__(
+        self,
+        *,
+        leaf_shape_index: collections.abc.Iterable[builtins.int] | None = ...,
+        instruction_name: builtins.str | None = ...,
+        shape_index: collections.abc.Iterable[builtins.int] | None = ...,
+    ) -> None: ...
+    def ClearField(
+        self,
+        field_name: typing.Literal[
+            "instruction_name", b"instruction_name", "leaf_shape_index", b"leaf_shape_index", "shape_index", b"shape_index"
+        ],
+    ) -> None: ...
+
+global___OriginalArrayProto = OriginalArrayProto
+
+@typing.final
+class OriginalValueProto(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    LEAVES_FIELD_NUMBER: builtins.int
+    @property
+    def leaves(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___OriginalArrayProto]: ...
+    def __init__(self, *, leaves: collections.abc.Iterable[global___OriginalArrayProto] | None = ...) -> None: ...
+    def ClearField(self, field_name: typing.Literal["leaves", b"leaves"]) -> None: ...
+
+global___OriginalValueProto = OriginalValueProto