VimPlug/typeshed
1
0 Fork 0
mirror of https://github.com/davidhalter/typeshed.git synced 2026-11-03 09:47:31 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Bump tensorflow to 2.17.* (#12512)

Browse Source
This commit is contained in:
Avasam authored and GitHub committed 2024-08-12 13:39:34 +02:00
1 parent dc0b63fd68
commit d3ca513ddb
15 files changed
+543 -128
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

stubs/tensorflow/METADATA.toml
+2 -2
View File
@@ -1,10 +1,10 @@
# Whenever you update version here, TENSORFLOW_VERSION should be updated
# in scripts/sync_tensorflow_protobuf_stubs.sh and vice-versa.
version = "2.16.*"
version = "2.17.*"
upstream_repository = "https://github.com/tensorflow/tensorflow"
# requires a version of numpy with a `py.typed` file
requires = ["numpy>=1.20", "types-protobuf", "types-requests"]
extra_description = "Partially generated using [mypy-protobuf==3.6.0](https://github.com/nipunn1313/mypy-protobuf/tree/v3.6.0) and libprotoc 25.1 on tensorflow==2.16.1 ."
extra_description = "Partially generated using [mypy-protobuf==3.6.0](https://github.com/nipunn1313/mypy-protobuf/tree/v3.6.0) and libprotoc 26.1 on tensorflow==2.17.0 ."
partial_stub = true
[tool.stubtest]
stubs/tensorflow/tensorflow/compiler/xla/service/hlo_pb2.pyi
+3 -3
View File
@@ -546,7 +546,7 @@ class HloInstructionProto(google.protobuf.message.Message):
"""
@property
def dot_sparsity(self) -> tensorflow.compiler.xla.xla_data_pb2.SparsityDescriptor:
def dot_sparsity(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[tensorflow.compiler.xla.xla_data_pb2.SparsityDescriptor]:
"""Sparsity descriptor for dot operation."""
def __init__(
@@ -622,9 +622,9 @@ class HloInstructionProto(google.protobuf.message.Message):
k: builtins.int | None = ...,
largest: builtins.bool | None = ...,
statistics_viz: tensorflow.compiler.xla.xla_data_pb2.StatisticsViz | None = ...,
dot_sparsity: tensorflow.compiler.xla.xla_data_pb2.SparsityDescriptor | None = ...,
dot_sparsity: collections.abc.Iterable[tensorflow.compiler.xla.xla_data_pb2.SparsityDescriptor] | None = ...,
) -> None: ...
def HasField(self, field_name: typing.Literal["cholesky_options", b"cholesky_options", "convolution_dimension_numbers", b"convolution_dimension_numbers", "cross_program_prefetch_index", b"cross_program_prefetch_index", "domain_entry_sharding", b"domain_entry_sharding", "domain_exit_sharding", b"domain_exit_sharding", "dot_dimension_numbers", b"dot_dimension_numbers", "dot_sparsity", b"dot_sparsity", "frontend_attributes", b"frontend_attributes", "gather_dimension_numbers", b"gather_dimension_numbers", "literal", b"literal", "metadata", b"metadata", "optional_cross_program_prefetch_index", b"optional_cross_program_prefetch_index", "outfeed_shape", b"outfeed_shape", "padding_config", b"padding_config", "parameter_replication", b"parameter_replication", "precision_config", b"precision_config", "scatter_dimension_numbers", b"scatter_dimension_numbers", "shape", b"shape", "sharding", b"sharding", "statistics_viz", b"statistics_viz", "triangular_solve_options", b"triangular_solve_options", "window", b"window"]) -> builtins.bool: ...
def HasField(self, field_name: typing.Literal["cholesky_options", b"cholesky_options", "convolution_dimension_numbers", b"convolution_dimension_numbers", "cross_program_prefetch_index", b"cross_program_prefetch_index", "domain_entry_sharding", b"domain_entry_sharding", "domain_exit_sharding", b"domain_exit_sharding", "dot_dimension_numbers", b"dot_dimension_numbers", "frontend_attributes", b"frontend_attributes", "gather_dimension_numbers", b"gather_dimension_numbers", "literal", b"literal", "metadata", b"metadata", "optional_cross_program_prefetch_index", b"optional_cross_program_prefetch_index", "outfeed_shape", b"outfeed_shape", "padding_config", b"padding_config", "parameter_replication", b"parameter_replication", "precision_config", b"precision_config", "scatter_dimension_numbers", b"scatter_dimension_numbers", "shape", b"shape", "sharding", b"sharding", "statistics_viz", b"statistics_viz", "triangular_solve_options", b"triangular_solve_options", "window", b"window"]) -> builtins.bool: ...
def ClearField(self, field_name: typing.Literal["all_reduce_id", b"all_reduce_id", "async_execution_thread", b"async_execution_thread", "backend_config", b"backend_config", "batch_group_count", b"batch_group_count", "called_computation_ids", b"called_computation_ids", "channel_id", b"channel_id", "cholesky_options", b"cholesky_options", "comparison_direction", b"comparison_direction", "comparison_type", b"comparison_type", "constrain_layout", b"constrain_layout", "control_predecessor_ids", b"control_predecessor_ids", "convolution_dimension_numbers", b"convolution_dimension_numbers", "cross_program_prefetch_index", b"cross_program_prefetch_index", "custom_call_api_version", b"custom_call_api_version", "custom_call_has_side_effect", b"custom_call_has_side_effect", "custom_call_schedule", b"custom_call_schedule", "custom_call_target", b"custom_call_target", "delta", b"delta", "dimensions", b"dimensions", "distribution", b"distribution", "domain_entry_sharding", b"domain_entry_sharding", "domain_exit_sharding", b"domain_exit_sharding", "dot_dimension_numbers", b"dot_dimension_numbers", "dot_sparsity", b"dot_sparsity", "dynamic_slice_sizes", b"dynamic_slice_sizes", "epsilon", b"epsilon", "exponent_bits", b"exponent_bits", "feature_group_count", b"feature_group_count", "feature_index", b"feature_index", "fft_length", b"fft_length", "fft_type", b"fft_type", "frontend_attributes", b"frontend_attributes", "fusion_kind", b"fusion_kind", "gather_dimension_numbers", b"gather_dimension_numbers", "gather_slice_sizes", b"gather_slice_sizes", "id", b"id", "indices_are_sorted", b"indices_are_sorted", "infeed_config", b"infeed_config", "is_cross_program_prefetch", b"is_cross_program_prefetch", "is_host_transfer", b"is_host_transfer", "is_stable", b"is_stable", "k", b"k", "largest", b"largest", "literal", b"literal", "mantissa_bits", b"mantissa_bits", "metadata", b"metadata", "name", b"name", "opcode", b"opcode", "operand_ids", b"operand_ids", "operand_shapes_with_layout", b"operand_shapes_with_layout", "optional_cross_program_prefetch_index", b"optional_cross_program_prefetch_index", "outfeed_config", b"outfeed_config", "outfeed_shape", b"outfeed_shape", "output_operand_aliasing", b"output_operand_aliasing", "padding_config", b"padding_config", "padding_type", b"padding_type", "parameter_number", b"parameter_number", "parameter_replication", b"parameter_replication", "precision_config", b"precision_config", "replica_groups", b"replica_groups", "rng_algorithm", b"rng_algorithm", "scatter_dimension_numbers", b"scatter_dimension_numbers", "shape", b"shape", "sharding", b"sharding", "slice_dimensions", b"slice_dimensions", "source_target_pairs", b"source_target_pairs", "statistics_viz", b"statistics_viz", "triangular_solve_options", b"triangular_solve_options", "tuple_index", b"tuple_index", "unique_indices", b"unique_indices", "use_global_device_ids", b"use_global_device_ids", "window", b"window"]) -> None: ...
def WhichOneof(self, oneof_group: typing.Literal["optional_cross_program_prefetch_index", b"optional_cross_program_prefetch_index"]) -> typing.Literal["cross_program_prefetch_index"] | None: ...
stubs/tensorflow/tensorflow/compiler/xla/service/metrics_pb2.pyi
+7 -1
View File
@@ -93,6 +93,7 @@ class CompilationLogEntry(google.protobuf.message.Message):
DURATION_FIELD_NUMBER: builtins.int
TASK_INDEX_FIELD_NUMBER: builtins.int
PASS_METRICS_FIELD_NUMBER: builtins.int
MODULE_IDS_FIELD_NUMBER: builtins.int
stage: global___CompilationLogEntry.CompilationStage.ValueType
"""Compilation stage recorded by this log entry."""
task_index: builtins.int
@@ -111,6 +112,10 @@ class CompilationLogEntry(google.protobuf.message.Message):
def pass_metrics(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___PassMetrics]:
"""Pass specific metrics."""
@property
def module_ids(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
"""IDs of modules on which the compilation stage was run."""
def __init__(
self,
*,
@@ -119,8 +124,9 @@ class CompilationLogEntry(google.protobuf.message.Message):
duration: google.protobuf.duration_pb2.Duration | None = ...,
task_index: builtins.int | None = ...,
pass_metrics: collections.abc.Iterable[global___PassMetrics] | None = ...,
module_ids: collections.abc.Iterable[builtins.int] | None = ...,
) -> None: ...
def HasField(self, field_name: typing.Literal["duration", b"duration", "timestamp", b"timestamp"]) -> builtins.bool: ...
def ClearField(self, field_name: typing.Literal["duration", b"duration", "pass_metrics", b"pass_metrics", "stage", b"stage", "task_index", b"task_index", "timestamp", b"timestamp"]) -> None: ...
def ClearField(self, field_name: typing.Literal["duration", b"duration", "module_ids", b"module_ids", "pass_metrics", b"pass_metrics", "stage", b"stage", "task_index", b"task_index", "timestamp", b"timestamp"]) -> None: ...
global___CompilationLogEntry = CompilationLogEntry
stubs/tensorflow/tensorflow/compiler/xla/xla_data_pb2.pyi
+164 -30
View File
@@ -44,14 +44,16 @@ class _PrimitiveTypeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._
"""Invalid primitive type to serve as default."""
PRED: _PrimitiveType.ValueType # 1
"""Predicates are two-state booleans."""
S4: _PrimitiveType.ValueType # 21
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
U4: _PrimitiveType.ValueType # 22
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
@@ -149,14 +151,16 @@ PRIMITIVE_TYPE_INVALID: PrimitiveType.ValueType # 0
"""Invalid primitive type to serve as default."""
PRED: PrimitiveType.ValueType # 1
"""Predicates are two-state booleans."""
S4: PrimitiveType.ValueType # 21
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
U4: PrimitiveType.ValueType # 22
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
@@ -508,8 +512,8 @@ global___PaddingConfig = PaddingConfig
@typing.final
class TileProto(google.protobuf.message.Message):
"""Describes a tile used in tiling-based layout. Refer to
g3doc/third_party/tensorflow/compiler/xla/g3doc/tiled_layout.md for
details about tiling-based layout.
g3doc/third_party/xla/docs/tiled_layout.md for details about tiling-based
layout.
"""
DESCRIPTOR: google.protobuf.descriptor.Descriptor
@@ -532,6 +536,33 @@ class TileProto(google.protobuf.message.Message):
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
@@ -560,6 +591,7 @@ class LayoutProto(google.protobuf.message.Message):
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
@@ -632,6 +664,12 @@ class LayoutProto(google.protobuf.message.Message):
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,
*,
@@ -647,9 +685,10 @@ class LayoutProto(google.protobuf.message.Message):
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", "tail_padding_alignment_in_elements", b"tail_padding_alignment_in_elements", "tiles", b"tiles"]) -> None: ...
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
@@ -815,8 +854,6 @@ class OpMetadata(google.protobuf.message.Message):
SOURCE_FILE_FIELD_NUMBER: builtins.int
SOURCE_LINE_FIELD_NUMBER: builtins.int
PROFILE_TYPE_FIELD_NUMBER: builtins.int
CREATION_PASS_ID_FIELD_NUMBER: builtins.int
LOGICAL_CREATION_PASS_ID_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
@@ -844,17 +881,6 @@ class OpMetadata(google.protobuf.message.Message):
e.g. it could be the file and line of user code that generated the op.
"""
source_line: builtins.int
creation_pass_id: builtins.int
"""HloPassMetadata.pass_id of the pass that created this HLO instruction
object. Should never be copied between HLO instructions. Zero if unset and
-1 if the instruction was created before HLO passes began.
"""
logical_creation_pass_id: builtins.int
"""HloPassMetadata.pass_id of the pass that created the logical functionality
that this HLO instruction represents. Should be copied between HLO
instructions that correspond across compilation passes. Zero if unset and
-1 if the instruction was created before HLO passes began.
"""
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
@@ -891,8 +917,6 @@ class OpMetadata(google.protobuf.message.Message):
source_file: builtins.str | None = ...,
source_line: builtins.int | None = ...,
profile_type: collections.abc.Iterable[global___ProfileType.ValueType] | None = ...,
creation_pass_id: builtins.int | None = ...,
logical_creation_pass_id: builtins.int | 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 = ...,
@@ -901,7 +925,7 @@ class OpMetadata(google.protobuf.message.Message):
stack_frame_id: builtins.int | None = ...,
) -> None: ...
def HasField(self, field_name: typing.Literal["profile_info", b"profile_info"]) -> builtins.bool: ...
def ClearField(self, field_name: typing.Literal["creation_pass_id", b"creation_pass_id", "deduplicated_name", b"deduplicated_name", "logical_creation_pass_id", b"logical_creation_pass_id", "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", "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: ...
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", "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
@@ -918,6 +942,7 @@ class ExecutionProfile(google.protobuf.message.Message):
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
@@ -944,6 +969,10 @@ class ExecutionProfile(google.protobuf.message.Message):
"""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,
*,
@@ -954,8 +983,9 @@ class ExecutionProfile(google.protobuf.message.Message):
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"]) -> 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
@@ -1139,9 +1169,11 @@ class LiteralProto(google.protobuf.message.Message):
SHAPE_FIELD_NUMBER: builtins.int
PREDS_FIELD_NUMBER: builtins.int
S2S_FIELD_NUMBER: builtins.int
S4S_FIELD_NUMBER: builtins.int
U4S_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
@@ -1162,9 +1194,11 @@ class LiteralProto(google.protobuf.message.Message):
F8E5M2FNUZS_FIELD_NUMBER: builtins.int
F8E4M3FNUZS_FIELD_NUMBER: builtins.int
SPARSE_INDICES_FIELD_NUMBER: builtins.int
s2s: builtins.bytes
s4s: builtins.bytes
u4s: 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"""
@@ -1204,16 +1238,18 @@ class LiteralProto(google.protobuf.message.Message):
def tuple_literals(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___LiteralProto]: ...
@property
def sparse_indices(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]:
"""Next = 26"""
"""Next = 28"""
def __init__(
self,
*,
shape: global___ShapeProto | None = ...,
preds: collections.abc.Iterable[builtins.bool] | None = ...,
s2s: builtins.bytes | None = ...,
s4s: builtins.bytes | None = ...,
u4s: 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 = ...,
@@ -1236,7 +1272,7 @@ class LiteralProto(google.protobuf.message.Message):
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", "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", "u32s", b"u32s", "u4s", b"u4s", "u64s", b"u64s", "u8s", b"u8s"]) -> None: ...
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
@@ -1992,15 +2028,113 @@ class PrecisionConfig(google.protobuf.message.Message):
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["operand_precision", b"operand_precision"]) -> None: ...
def ClearField(self, field_name: typing.Literal["algorithm", b"algorithm", "operand_precision", b"operand_precision"]) -> None: ...
global___PrecisionConfig = PrecisionConfig
stubs/tensorflow/tensorflow/compiler/xla/xla_pb2.pyi
+239 -82
View File
@@ -137,6 +137,33 @@ class DebugOptions(google.protobuf.message.Message):
STEP_MARK_NONE: DebugOptions.StepMarkerLocation.ValueType # 2
"""No step marker generated."""
class _CollectiveOpType:
ValueType = typing.NewType("ValueType", builtins.int)
V: typing_extensions.TypeAlias = ValueType
class _CollectiveOpTypeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[DebugOptions._CollectiveOpType.ValueType], builtins.type):
DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
NOOP: DebugOptions._CollectiveOpType.ValueType # 0
ALLREDUCE: DebugOptions._CollectiveOpType.ValueType # 1
ALLGATHER: DebugOptions._CollectiveOpType.ValueType # 2
REDUCESCATTER: DebugOptions._CollectiveOpType.ValueType # 3
COLLECTIVEBROADCAST: DebugOptions._CollectiveOpType.ValueType # 4
ALLTOALL: DebugOptions._CollectiveOpType.ValueType # 5
COLLECTIVEPERMUTE: DebugOptions._CollectiveOpType.ValueType # 6
class CollectiveOpType(_CollectiveOpType, metaclass=_CollectiveOpTypeEnumTypeWrapper):
"""Enum to define all collective ops
that xla supports.
"""
NOOP: DebugOptions.CollectiveOpType.ValueType # 0
ALLREDUCE: DebugOptions.CollectiveOpType.ValueType # 1
ALLGATHER: DebugOptions.CollectiveOpType.ValueType # 2
REDUCESCATTER: DebugOptions.CollectiveOpType.ValueType # 3
COLLECTIVEBROADCAST: DebugOptions.CollectiveOpType.ValueType # 4
ALLTOALL: DebugOptions.CollectiveOpType.ValueType # 5
COLLECTIVEPERMUTE: DebugOptions.CollectiveOpType.ValueType # 6
class _CommandBufferCmdType:
ValueType = typing.NewType("ValueType", builtins.int)
V: typing_extensions.TypeAlias = ValueType
@@ -150,11 +177,12 @@ class DebugOptions(google.protobuf.message.Message):
COLLECTIVES: DebugOptions._CommandBufferCmdType.ValueType # 4
CONDITIONALS: DebugOptions._CommandBufferCmdType.ValueType # 5
CUSTOM_CALL: DebugOptions._CommandBufferCmdType.ValueType # 6
CUBLASLT: DebugOptions._CommandBufferCmdType.ValueType # 7
class CommandBufferCmdType(_CommandBufferCmdType, metaclass=_CommandBufferCmdTypeEnumTypeWrapper):
"""Commands are categorized into 5 types:
FUSION represents regular fusion kernels.
CUBLAS, CUDNN, and COLLECTIVES represent library calls.
CUBLAS/CUBLASLT, CUDNN, and COLLECTIVES represent library calls.
CONDITIONALS represents control flow.
"""
@@ -165,6 +193,7 @@ class DebugOptions(google.protobuf.message.Message):
COLLECTIVES: DebugOptions.CommandBufferCmdType.ValueType # 4
CONDITIONALS: DebugOptions.CommandBufferCmdType.ValueType # 5
CUSTOM_CALL: DebugOptions.CommandBufferCmdType.ValueType # 6
CUBLASLT: DebugOptions.CommandBufferCmdType.ValueType # 7
class _PartitioningAlgorithm:
ValueType = typing.NewType("ValueType", builtins.int)
@@ -183,6 +212,29 @@ class DebugOptions(google.protobuf.message.Message):
PARTITIONING_ALGORITHM_EXP1: DebugOptions.PartitioningAlgorithm.ValueType # 2
PARTITIONING_ALGORITHM_EXP2: DebugOptions.PartitioningAlgorithm.ValueType # 3
class _WhileLoopUnrolling:
ValueType = typing.NewType("ValueType", builtins.int)
V: typing_extensions.TypeAlias = ValueType
class _WhileLoopUnrollingEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[DebugOptions._WhileLoopUnrolling.ValueType], builtins.type):
DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
WHILE_LOOP_UNROLLING_NO_UNROLL: DebugOptions._WhileLoopUnrolling.ValueType # 0
WHILE_LOOP_UNROLLING_DOUBLE_BUFFER: DebugOptions._WhileLoopUnrolling.ValueType # 1
"""Has the same effect as setting
`xla_gpu_enable_while_loop_double_buffering`.
"""
WHILE_LOOP_UNROLLING_FULL_UNROLL: DebugOptions._WhileLoopUnrolling.ValueType # 2
"""Enables full loop unrolling using the same strategy as `DOUBLE_BUFFER`."""
class WhileLoopUnrolling(_WhileLoopUnrolling, metaclass=_WhileLoopUnrollingEnumTypeWrapper): ...
WHILE_LOOP_UNROLLING_NO_UNROLL: DebugOptions.WhileLoopUnrolling.ValueType # 0
WHILE_LOOP_UNROLLING_DOUBLE_BUFFER: DebugOptions.WhileLoopUnrolling.ValueType # 1
"""Has the same effect as setting
`xla_gpu_enable_while_loop_double_buffering`.
"""
WHILE_LOOP_UNROLLING_FULL_UNROLL: DebugOptions.WhileLoopUnrolling.ValueType # 2
"""Enables full loop unrolling using the same strategy as `DOUBLE_BUFFER`."""
@typing.final
class XlaBackendExtraOptionsEntry(google.protobuf.message.Message):
DESCRIPTOR: google.protobuf.descriptor.Descriptor
@@ -218,7 +270,7 @@ class DebugOptions(google.protobuf.message.Message):
XLA_TEST_ALL_INPUT_LAYOUTS_FIELD_NUMBER: builtins.int
XLA_HLO_GRAPH_SHARDING_COLOR_FIELD_NUMBER: builtins.int
XLA_CPU_USE_MKL_DNN_FIELD_NUMBER: builtins.int
XLA_CPU_USE_XLA_RUNTIME_FIELD_NUMBER: builtins.int
XLA_CPU_USE_THUNK_RUNTIME_FIELD_NUMBER: builtins.int
XLA_CPU_ENABLE_FAST_MATH_FIELD_NUMBER: builtins.int
XLA_CPU_FAST_MATH_HONOR_NANS_FIELD_NUMBER: builtins.int
XLA_CPU_FAST_MATH_HONOR_INFS_FIELD_NUMBER: builtins.int
@@ -226,7 +278,6 @@ class DebugOptions(google.protobuf.message.Message):
XLA_CPU_FAST_MATH_HONOR_FUNCTIONS_FIELD_NUMBER: builtins.int
XLA_CPU_ENABLE_FAST_MIN_MAX_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_FAST_MIN_MAX_FIELD_NUMBER: builtins.int
XLA_CPU_SPARSE_CUDA_THREADS_FIELD_NUMBER: builtins.int
XLA_ALLOW_EXCESS_PRECISION_FIELD_NUMBER: builtins.int
XLA_GPU_CRASH_ON_VERIFICATION_FAILURES_FIELD_NUMBER: builtins.int
XLA_GPU_AUTOTUNE_LEVEL_FIELD_NUMBER: builtins.int
@@ -269,22 +320,19 @@ class DebugOptions(google.protobuf.message.Message):
XLA_GPU_ENABLE_LLVM_MODULE_COMPILATION_PARALLELISM_FIELD_NUMBER: builtins.int
XLA_GPU_DETERMINISTIC_OPS_FIELD_NUMBER: builtins.int
XLA_GPU_LLVM_IR_FILE_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_ASYNC_COLLECTIVES_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_ASYNC_ALL_REDUCE_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_ASYNC_COLLECTIVE_PERMUTE_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_ASYNC_ALL_GATHER_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_ASYNC_REDUCE_SCATTER_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_ASYNC_ALL_TO_ALL_FIELD_NUMBER: builtins.int
XLA_GPU_DISABLE_ASYNC_COLLECTIVES_FIELD_NUMBER: builtins.int
XLA_GPU_ALL_REDUCE_COMBINE_THRESHOLD_BYTES_FIELD_NUMBER: builtins.int
XLA_GPU_ALL_GATHER_COMBINE_THRESHOLD_BYTES_FIELD_NUMBER: builtins.int
XLA_GPU_REDUCE_SCATTER_COMBINE_THRESHOLD_BYTES_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_ALL_GATHER_COMBINE_BY_DIM_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_REDUCE_SCATTER_COMBINE_BY_DIM_FIELD_NUMBER: builtins.int
XLA_GPU_ALL_REDUCE_CONTIGUOUS_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_ALL_REDUCE_SPLITTER_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_REASSOCIATION_FOR_CONVERTED_AR_FIELD_NUMBER: builtins.int
XLA_GPU_ALL_REDUCE_BLUECONNECT_NUM_DEVICES_PER_HOST_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_WHILE_LOOP_REDUCE_SCATTER_CODE_MOTION_FIELD_NUMBER: builtins.int
XLA_GPU_COLLECTIVE_INFLATION_FACTOR_FIELD_NUMBER: builtins.int
XLA_LLVM_FORCE_INLINE_BEFORE_SPLIT_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_CUDNN_FRONTEND_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_CUDNN_FMHA_FIELD_NUMBER: builtins.int
XLA_GPU_FUSED_ATTENTION_USE_CUDNN_RNG_FIELD_NUMBER: builtins.int
@@ -292,7 +340,6 @@ class DebugOptions(google.protobuf.message.Message):
XLA_DUMP_DISABLE_METADATA_FIELD_NUMBER: builtins.int
XLA_DUMP_HLO_PIPELINE_RE_FIELD_NUMBER: builtins.int
XLA_GPU_STRICT_CONV_ALGORITHM_PICKER_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_XLA_RUNTIME_EXECUTABLE_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_CUSTOM_FUSIONS_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_CUSTOM_FUSIONS_RE_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_ADDRESS_COMPUTATION_FUSION_FIELD_NUMBER: builtins.int
@@ -300,14 +347,10 @@ class DebugOptions(google.protobuf.message.Message):
XLA_GPU_ENABLE_SHARED_CONSTANTS_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_CUBLASLT_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_COMMAND_BUFFER_FIELD_NUMBER: builtins.int
XLA_GPU_GRAPH_NUM_RUNS_TO_INSTANTIATE_FIELD_NUMBER: builtins.int
XLA_GPU_GRAPH_MIN_GRAPH_SIZE_FIELD_NUMBER: builtins.int
XLA_GPU_GRAPH_ENABLE_CONCURRENT_REGION_FIELD_NUMBER: builtins.int
XLA_GPU_GRAPH_EVICTION_TIMEOUT_SECONDS_FIELD_NUMBER: builtins.int
XLA_GPU_REDZONE_SCRATCH_MAX_MEGABYTES_FIELD_NUMBER: builtins.int
XLA_GPU_REDZONE_PADDING_BYTES_FIELD_NUMBER: builtins.int
XLA_GPU_SIMPLIFY_ALL_FP_CONVERSIONS_FIELD_NUMBER: builtins.int
XLA_GPU_NORMALIZE_LAYOUTS_FIELD_NUMBER: builtins.int
XLA_CPU_USE_ACL_FIELD_NUMBER: builtins.int
XLA_CPU_STRICT_DOT_CONV_MATH_FIELD_NUMBER: builtins.int
XLA_GPU_USE_RUNTIME_FUSION_FIELD_NUMBER: builtins.int
@@ -338,6 +381,7 @@ class DebugOptions(google.protobuf.message.Message):
XLA_GPU_EXHAUSTIVE_TILING_SEARCH_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_TRITON_SOFTMAX_FUSION_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_PRIORITY_FUSION_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_TRITON_SOFTMAX_PRIORITY_FUSION_FIELD_NUMBER: builtins.int
XLA_GPU_DUMP_AUTOTUNE_RESULTS_TO_FIELD_NUMBER: builtins.int
XLA_GPU_LOAD_AUTOTUNE_RESULTS_FROM_FIELD_NUMBER: builtins.int
XLA_GPU_TARGET_CONFIG_FILENAME_FIELD_NUMBER: builtins.int
@@ -345,7 +389,8 @@ class DebugOptions(google.protobuf.message.Message):
XLA_GPU_AUTO_SPMD_PARTITIONING_MEMORY_BUDGET_RATIO_FIELD_NUMBER: builtins.int
XLA_GPU_TRITON_GEMM_DISABLE_REDUCED_PRECISION_REDUCTION_FIELD_NUMBER: builtins.int
XLA_GPU_TRITON_FUSION_LEVEL_FIELD_NUMBER: builtins.int
XLA_GPU_DUMP_AUTOTUNED_TRITON_FUSIONS_FIELD_NUMBER: builtins.int
XLA_GPU_DUMP_AUTOTUNED_GEMM_FUSIONS_FIELD_NUMBER: builtins.int
XLA_GPU_OVERRIDE_GEMM_AUTOTUNER_FIELD_NUMBER: builtins.int
XLA_GPU_COPY_INSERTION_USE_REGION_ANALYSIS_FIELD_NUMBER: builtins.int
XLA_GPU_COLLECT_COST_MODEL_STATS_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_SPLIT_K_AUTOTUNING_FIELD_NUMBER: builtins.int
@@ -354,6 +399,7 @@ class DebugOptions(google.protobuf.message.Message):
XLA_GPU_MOCK_CUSTOM_CALLS_FIELD_NUMBER: builtins.int
XLA_GPU_CUBLAS_FALLBACK_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_WHILE_LOOP_DOUBLE_BUFFERING_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_WHILE_LOOP_UNROLLING_FIELD_NUMBER: builtins.int
XLA_GPU_ENSURE_MINOR_DOT_CONTRACTION_DIMS_FIELD_NUMBER: builtins.int
XLA_GPU_FILTER_KERNELS_SPILLING_REGISTERS_ON_AUTOTUNING_FIELD_NUMBER: builtins.int
XLA_DEBUG_BUFFER_ASSIGNMENT_SHOW_MAX_FIELD_NUMBER: builtins.int
@@ -362,8 +408,29 @@ class DebugOptions(google.protobuf.message.Message):
XLA_GPU_THRESHOLD_FOR_WINDOWED_EINSUM_MIB_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_TRITON_HOPPER_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_NCCL_USER_BUFFERS_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_NCCL_COMM_SPLITTING_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_NCCL_PER_STREAM_COMMS_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_LIBNVPTXCOMPILER_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_DOT_STRENGTH_REDUCTION_FIELD_NUMBER: builtins.int
XLA_GPU_MULTI_STREAMED_WINDOWED_EINSUM_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_BF16_6WAY_GEMM_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_BF16_3WAY_GEMM_FIELD_NUMBER: builtins.int
XLA_GPU_NCCL_COLLECTIVE_MAX_NCHANNELS_FIELD_NUMBER: builtins.int
XLA_GPU_NCCL_P2P_MAX_NCHANNELS_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_MLIR_EMITTERS_FIELD_NUMBER: builtins.int
XLA_GPU_MAX_MLIR_KERNELS_FIELD_NUMBER: builtins.int
XLA_GPU_SKIP_MLIR_KERNELS_FIELD_NUMBER: builtins.int
XLA_GPU_GEMM_REWRITE_SIZE_THRESHOLD_FIELD_NUMBER: builtins.int
XLA_GPU_REQUIRE_COMPLETE_AOT_AUTOTUNE_RESULTS_FIELD_NUMBER: builtins.int
XLA_GPU_CUDNN_GEMM_FUSION_LEVEL_FIELD_NUMBER: builtins.int
XLA_GPU_USE_MEMCPY_LOCAL_P2P_FIELD_NUMBER: builtins.int
XLA_GPU_AUTOTUNE_MAX_SOLUTIONS_FIELD_NUMBER: builtins.int
XLA_DUMP_LARGE_CONSTANTS_FIELD_NUMBER: builtins.int
XLA_GPU_VERIFY_TRITON_FUSION_NUMERICS_FIELD_NUMBER: builtins.int
XLA_GPU_DUMP_AUTOTUNE_LOGS_TO_FIELD_NUMBER: builtins.int
XLA_REDUCE_WINDOW_REWRITE_BASE_LENGTH_FIELD_NUMBER: builtins.int
XLA_GPU_ENABLE_HOST_MEMORY_OFFLOADING_FIELD_NUMBER: builtins.int
XLA_GPU_EXCLUDE_NONDETERMINISTIC_OPS_FIELD_NUMBER: builtins.int
XLA_BACKEND_EXTRA_OPTIONS_FIELD_NUMBER: builtins.int
xla_hlo_graph_addresses: builtins.bool
"""Show addresses of HLO ops in graph dump."""
@@ -427,8 +494,7 @@ class DebugOptions(google.protobuf.message.Message):
"""
xla_cpu_use_mkl_dnn: builtins.bool
"""Generate calls to MKL-DNN in the CPU backend."""
xla_cpu_use_xla_runtime: builtins.bool
"""Enable XLA Runtime in the CPU backend."""
xla_cpu_use_thunk_runtime: builtins.bool
xla_cpu_enable_fast_math: builtins.bool
"""When true, "unsafe" mathematical optimizations are enabled. These
transformations include but are not limited to:
@@ -475,15 +541,14 @@ class DebugOptions(google.protobuf.message.Message):
Note, this does not correspond to the exact same behavior as the cpu flag
above!
"""
xla_cpu_sparse_cuda_threads: builtins.int
"""Defines the number of CUDA threads that can be used to accelerate
a sparse computation compiled for the XLA Runtime and CPU backend.
By default (value 0), no acceleration is used. Otherwise, this
many threads may be used to accelerate sparse operations, typically
useful when accelerating structured sparsity.
"""
xla_allow_excess_precision: builtins.bool
"""Allows xla to increase the output precision of floating point operations."""
"""Allows xla to increase the output precision of floating point operations
and all floating-point conversions to be simplified, including those
that affect the numerics. The `FloatNormalization` pass inserts many
`f32 -> bf16 -> f32` conversion pairs. These are not removed by the
`AlgebraicSimplifier`, as that will only simplify conversions that are
no-ops, e.g. `bf16 -> f32 -> bf16`. Removing these improves accuracy.
"""
xla_gpu_crash_on_verification_failures: builtins.bool
"""Crashes the program when any kind of verification fails, instead of just
logging the failures. One example is cross checking of convolution results
@@ -612,17 +677,10 @@ class DebugOptions(google.protobuf.message.Message):
"""
xla_gpu_enable_llvm_module_compilation_parallelism: builtins.bool
xla_gpu_deterministic_ops: builtins.bool
"""Guarantees run-to-run determinism. At present, the HLO ops Scatter and
SelectAndScatter do not have deterministic XLA:GPU implementations.
Compilation errors out if these ops are encountered.
"""Guarantees run-to-run determinism.
This flag implies --xla_gpu_exclude_nondeterministic_ops and in addition
disables autotuning.
"""
xla_gpu_enable_async_collectives: builtins.bool
"""Convert synchronous collective ops into asynchronous."""
xla_gpu_enable_async_all_reduce: builtins.bool
xla_gpu_enable_async_collective_permute: builtins.bool
xla_gpu_enable_async_all_gather: builtins.bool
xla_gpu_enable_async_reduce_scatter: builtins.bool
xla_gpu_enable_async_all_to_all: builtins.bool
xla_gpu_all_reduce_combine_threshold_bytes: builtins.int
"""Size threshold (in bytes) for the GPU collective combiners."""
xla_gpu_all_gather_combine_threshold_bytes: builtins.int
@@ -634,6 +692,8 @@ class DebugOptions(google.protobuf.message.Message):
xla_gpu_enable_reduce_scatter_combine_by_dim: builtins.bool
xla_gpu_all_reduce_contiguous: builtins.bool
"""Combine GPU all-reduces into a single operation over a contiguous buffer."""
xla_gpu_enable_all_reduce_splitter: builtins.bool
"""Split AR across all partitions into logical RS followed by AR."""
xla_gpu_enable_reassociation_for_converted_ar: builtins.bool
"""Enable allreduce reassociation on allreduces that are converted to a wider
type. The resulting allreduce will be promoted to a wider-typed allreduce.
@@ -649,6 +709,10 @@ class DebugOptions(google.protobuf.message.Message):
"""Enable hoisting of reduce-scatter out of while loops."""
xla_gpu_collective_inflation_factor: builtins.int
"""Inflate collective cost by running each collective multiple times."""
xla_llvm_force_inline_before_split: builtins.bool
"""Whether to force inline before llvm module split to get a more balanced
splits for parallel compilation.
"""
xla_gpu_enable_cudnn_frontend: builtins.bool
"""Whether to use the cuDNN frontend API for convolutions when possible."""
xla_gpu_enable_cudnn_fmha: builtins.bool
@@ -666,8 +730,6 @@ class DebugOptions(google.protobuf.message.Message):
"""If true, abort immediately when conv algorithm picker fails, rather than
logging a warning and proceeding with fallback.
"""
xla_gpu_enable_xla_runtime_executable: builtins.bool
"""If true, use XLA runtime for XLA:GPU backend."""
xla_gpu_enable_custom_fusions: builtins.bool
"""If true, XLA will try to pattern match subgraphs of HLO operations into
custom fusions registered in the current process (pre-compiled hand written
@@ -678,7 +740,9 @@ class DebugOptions(google.protobuf.message.Message):
if `xla_gpu_enable_custom_fusion` set to true.
"""
xla_gpu_enable_address_computation_fusion: builtins.bool
"""If true, use XLA runtime for XLA:GPU backend."""
"""Enables address computation fusion to optimize dynamic-slice and
dynamic-update-slice operations around library calls.
"""
xla_gpu_nccl_termination_timeout_seconds: builtins.int
"""Timeout in seconds before terminating jobs that are stuck in a NCCL
Rendezvous. Negative value disables the timeout and will not terminate.
@@ -689,11 +753,6 @@ class DebugOptions(google.protobuf.message.Message):
"""
xla_gpu_enable_cublaslt: builtins.bool
"""Whether to use cuBLASLt for GEMMs on GPUs."""
xla_gpu_graph_num_runs_to_instantiate: builtins.int
"""Only instantiates a GPU graph after the captured function execution count
reaches the threshold. This constant is a heuristic to avoid creating a
large number of CUDA graph instances in memory.
"""
xla_gpu_graph_min_graph_size: builtins.int
"""This number determines how many moved instructions like fusion kernels are
required for a region to be captured as a function to be launched as a GPU
@@ -701,11 +760,6 @@ class DebugOptions(google.protobuf.message.Message):
"""
xla_gpu_graph_enable_concurrent_region: builtins.bool
"""Identify concurrent regions in GPU graphs and execute them concurrently."""
xla_gpu_graph_eviction_timeout_seconds: builtins.int
"""Timeout in seconds to evict instantiated Gpu graphs from device. When XLA
instantiates new Gpu graphs, it evicts graphs that were not recently
executed to free space on device.
"""
xla_gpu_redzone_scratch_max_megabytes: builtins.int
"""Size threshold (in megabytes) for the GPU redzone scratch allocator."""
xla_gpu_redzone_padding_bytes: builtins.int
@@ -718,19 +772,6 @@ class DebugOptions(google.protobuf.message.Message):
fused cudnn conv has up to 6 total buffers (4 inputs, 1 output, and 1
scratch), so this can be multiplied by quite a lot.
"""
xla_gpu_simplify_all_fp_conversions: builtins.bool
"""Allows all floating-point conversions to be simplified, including those
that affect the numerics. The `FloatNormalization` pass inserts many
`f32 -> bf16 -> f32` conversion pairs. These are not removed by the
`AlgebraicSimplifier`, as that will only simplify conversions that are
no-ops, e.g. `bf16 -> f32 -> bf16`. Removing these improves accuracy.
"""
xla_gpu_normalize_layouts: builtins.bool
"""An experimental option to force all layouts present in the
after-optimizations HLO to be descending, e.g.
ShapeUtil::MakeShapeWithDescendingLayout is an identity on all
instructions.
"""
xla_cpu_use_acl: builtins.bool
"""Generate calls to Arm Compute Library in the CPU backend."""
xla_cpu_strict_dot_conv_math: builtins.bool
@@ -782,6 +823,7 @@ class DebugOptions(google.protobuf.message.Message):
xla_gpu_exhaustive_tiling_search: builtins.bool
xla_gpu_enable_triton_softmax_fusion: builtins.bool
xla_gpu_enable_priority_fusion: builtins.bool
xla_gpu_enable_triton_softmax_priority_fusion: builtins.bool
xla_gpu_dump_autotune_results_to: builtins.str
"""File to write autotune results to. It will be a binary file unless the name
ends with .txt or .textproto. Warning: The results are written at every
@@ -808,7 +850,8 @@ class DebugOptions(google.protobuf.message.Message):
"""
xla_gpu_triton_gemm_disable_reduced_precision_reduction: builtins.bool
xla_gpu_triton_fusion_level: builtins.int
xla_gpu_dump_autotuned_triton_fusions: builtins.bool
xla_gpu_dump_autotuned_gemm_fusions: builtins.bool
xla_gpu_override_gemm_autotuner: builtins.str
xla_gpu_copy_insertion_use_region_analysis: builtins.bool
xla_gpu_collect_cost_model_stats: builtins.bool
"""If true, each fusion instruction will have a cost model runtime estimate in
@@ -827,6 +870,8 @@ class DebugOptions(google.protobuf.message.Message):
"""
xla_gpu_enable_while_loop_double_buffering: builtins.bool
"""Enable double buffering for loops."""
xla_gpu_enable_while_loop_unrolling: global___DebugOptions.WhileLoopUnrolling.ValueType
"""Determine the while loop unrolling scheme."""
xla_gpu_ensure_minor_dot_contraction_dims: builtins.bool
"""Change the layout of the second triton dot operand to be column major.
Only works for (bf16 x bf16) -> bf16.
@@ -844,9 +889,84 @@ class DebugOptions(google.protobuf.message.Message):
"""Enables currently disabled features within Triton for Hopper."""
xla_gpu_enable_nccl_user_buffers: builtins.bool
"""Enable NCCL user buffers."""
xla_gpu_enable_nccl_comm_splitting: builtins.bool
"""Enable NCCL communicator splitting."""
xla_gpu_enable_nccl_per_stream_comms: builtins.bool
"""Enable NCCL per stream communicators."""
xla_gpu_enable_libnvptxcompiler: builtins.bool
"""If enabled, uses the libnvptxcompiler library to compile PTX to cuBIN."""
xla_gpu_enable_dot_strength_reduction: builtins.bool
xla_gpu_multi_streamed_windowed_einsum: builtins.bool
"""Whether to use multiple compute streams to run windowed einsum."""
xla_gpu_enable_bf16_6way_gemm: builtins.bool
"""If enabled, uses bf16_6way gemm to compute F32 gemm."""
xla_gpu_enable_bf16_3way_gemm: builtins.bool
"""If enabled, uses bf16_3way gemm to compute F32 gemm."""
xla_gpu_nccl_collective_max_nchannels: builtins.int
"""Specify the maximum number of channels(SMs) NCCL
will use for collective operations.
"""
xla_gpu_nccl_p2p_max_nchannels: builtins.int
"""Specify the maximum number of channels(SMs) NCCL
will use for p2p operations.
"""
xla_gpu_enable_mlir_emitters: builtins.bool
xla_gpu_max_mlir_kernels: builtins.int
"""The maximum number of kernels to emit with MLIR. Unlimited if 0."""
xla_gpu_skip_mlir_kernels: builtins.int
"""The number of initial kernels to not emit with MLIR. Only supported kernels
are counted.
"""
xla_gpu_gemm_rewrite_size_threshold: builtins.int
"""Threshold to rewrite matmul to cuBLAS or Triton (minumum combined number of
elements of both matrices in non-batch dimensions to be considered for a
rewrite).
"""
xla_gpu_require_complete_aot_autotune_results: builtins.bool
"""If true, will require complete AOT autotuning results; in the case of
missing AOT result, the model will not be compiled or executed, a
`NotFound` error will be returned.
"""
xla_gpu_cudnn_gemm_fusion_level: builtins.int
"""Let GEMM fusion autotuning probe cuDNN as a backend.
Current levels:
0: Disabled.
1: Fusions of GEMM, elementwise, transpose/reshape operations.
2: + Broadcasts.
3: + Nontrivial noncontracting dimension reshapes/transposes.
"""
xla_gpu_use_memcpy_local_p2p: builtins.bool
"""This instructs the runtime whether to use
memcpy for p2p communication when source and
target are located within a node(nvlink).
"""
xla_gpu_autotune_max_solutions: builtins.int
"""If non-zero, limits the number of solutions to be used by GEMM autotuner.
This might be useful if underlying math library returns too many GEMM
solutions.
"""
xla_dump_large_constants: builtins.bool
"""If true, large constants will be printed out when dumping HLOs."""
xla_gpu_verify_triton_fusion_numerics: builtins.bool
"""If true, will verify that the numerical results of Triton fusions match
the results of regular emitters.
"""
xla_gpu_dump_autotune_logs_to: builtins.str
"""File to write autotune logs to. It will stored in txt format."""
xla_reduce_window_rewrite_base_length: builtins.int
"""Base length to rewrite the reduce window to, no rewrite if set to 0."""
xla_gpu_enable_host_memory_offloading: builtins.bool
"""If true, will enable host memory offloading on a device."""
xla_gpu_exclude_nondeterministic_ops: builtins.bool
"""Excludes non-deterministic ops from compiled executables.
Unlike --xla_gpu_deterministic_ops does not disable autotuning - the
compilation itself can be non-deterministic.
At present, the HLO op SelectAndScatter does not have a
deterministic XLA:GPU implementation.
Compilation errors out if SelectAndScatter is encountered.
Scatter ops can non-deterministic by default; these get converted to
a deterministic implementation.
"""
@property
def xla_disable_hlo_passes(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]:
"""List of HLO passes to disable/enable. These names must exactly match the
@@ -866,13 +986,15 @@ class DebugOptions(google.protobuf.message.Message):
def xla_gpu_llvm_ir_file(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]:
"""Paths to files with LLVM code."""
@property
def xla_gpu_disable_async_collectives(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[global___DebugOptions.CollectiveOpType.ValueType]: ...
@property
def xla_gpu_enable_command_buffer(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[global___DebugOptions.CommandBufferCmdType.ValueType]:
"""Determine the types of commands that are recorded into command buffers."""
@property
def xla_backend_extra_options(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]:
"""Next id: 271
"""Next id: 300
Extra options to pass to the compilation backend (e.g. LLVM); specific
interpretation of these values is left to the backend.
@@ -900,7 +1022,7 @@ class DebugOptions(google.protobuf.message.Message):
xla_test_all_input_layouts: builtins.bool | None = ...,
xla_hlo_graph_sharding_color: builtins.bool | None = ...,
xla_cpu_use_mkl_dnn: builtins.bool | None = ...,
xla_cpu_use_xla_runtime: builtins.bool | None = ...,
xla_cpu_use_thunk_runtime: builtins.bool | None = ...,
xla_cpu_enable_fast_math: builtins.bool | None = ...,
xla_cpu_fast_math_honor_nans: builtins.bool | None = ...,
xla_cpu_fast_math_honor_infs: builtins.bool | None = ...,
@@ -908,7 +1030,6 @@ class DebugOptions(google.protobuf.message.Message):
xla_cpu_fast_math_honor_functions: builtins.bool | None = ...,
xla_cpu_enable_fast_min_max: builtins.bool | None = ...,
xla_gpu_enable_fast_min_max: builtins.bool | None = ...,
xla_cpu_sparse_cuda_threads: builtins.int | None = ...,
xla_allow_excess_precision: builtins.bool | None = ...,
xla_gpu_crash_on_verification_failures: builtins.bool | None = ...,
xla_gpu_autotune_level: builtins.int | None = ...,
@@ -951,22 +1072,19 @@ class DebugOptions(google.protobuf.message.Message):
xla_gpu_enable_llvm_module_compilation_parallelism: builtins.bool | None = ...,
xla_gpu_deterministic_ops: builtins.bool | None = ...,
xla_gpu_llvm_ir_file: collections.abc.Iterable[builtins.str] | None = ...,
xla_gpu_enable_async_collectives: builtins.bool | None = ...,
xla_gpu_enable_async_all_reduce: builtins.bool | None = ...,
xla_gpu_enable_async_collective_permute: builtins.bool | None = ...,
xla_gpu_enable_async_all_gather: builtins.bool | None = ...,
xla_gpu_enable_async_reduce_scatter: builtins.bool | None = ...,
xla_gpu_enable_async_all_to_all: builtins.bool | None = ...,
xla_gpu_disable_async_collectives: collections.abc.Iterable[global___DebugOptions.CollectiveOpType.ValueType] | None = ...,
xla_gpu_all_reduce_combine_threshold_bytes: builtins.int | None = ...,
xla_gpu_all_gather_combine_threshold_bytes: builtins.int | None = ...,
xla_gpu_reduce_scatter_combine_threshold_bytes: builtins.int | None = ...,
xla_gpu_enable_all_gather_combine_by_dim: builtins.bool | None = ...,
xla_gpu_enable_reduce_scatter_combine_by_dim: builtins.bool | None = ...,
xla_gpu_all_reduce_contiguous: builtins.bool | None = ...,
xla_gpu_enable_all_reduce_splitter: builtins.bool | None = ...,
xla_gpu_enable_reassociation_for_converted_ar: builtins.bool | None = ...,
xla_gpu_all_reduce_blueconnect_num_devices_per_host: builtins.int | None = ...,
xla_gpu_enable_while_loop_reduce_scatter_code_motion: builtins.bool | None = ...,
xla_gpu_collective_inflation_factor: builtins.int | None = ...,
xla_llvm_force_inline_before_split: builtins.bool | None = ...,
xla_gpu_enable_cudnn_frontend: builtins.bool | None = ...,
xla_gpu_enable_cudnn_fmha: builtins.bool | None = ...,
xla_gpu_fused_attention_use_cudnn_rng: builtins.bool | None = ...,
@@ -974,7 +1092,6 @@ class DebugOptions(google.protobuf.message.Message):
xla_dump_disable_metadata: builtins.bool | None = ...,
xla_dump_hlo_pipeline_re: builtins.str | None = ...,
xla_gpu_strict_conv_algorithm_picker: builtins.bool | None = ...,
xla_gpu_enable_xla_runtime_executable: builtins.bool | None = ...,
xla_gpu_enable_custom_fusions: builtins.bool | None = ...,
xla_gpu_enable_custom_fusions_re: builtins.str | None = ...,
xla_gpu_enable_address_computation_fusion: builtins.bool | None = ...,
@@ -982,14 +1099,10 @@ class DebugOptions(google.protobuf.message.Message):
xla_gpu_enable_shared_constants: builtins.bool | None = ...,
xla_gpu_enable_cublaslt: builtins.bool | None = ...,
xla_gpu_enable_command_buffer: collections.abc.Iterable[global___DebugOptions.CommandBufferCmdType.ValueType] | None = ...,
xla_gpu_graph_num_runs_to_instantiate: builtins.int | None = ...,
xla_gpu_graph_min_graph_size: builtins.int | None = ...,
xla_gpu_graph_enable_concurrent_region: builtins.bool | None = ...,
xla_gpu_graph_eviction_timeout_seconds: builtins.int | None = ...,
xla_gpu_redzone_scratch_max_megabytes: builtins.int | None = ...,
xla_gpu_redzone_padding_bytes: builtins.int | None = ...,
xla_gpu_simplify_all_fp_conversions: builtins.bool | None = ...,
xla_gpu_normalize_layouts: builtins.bool | None = ...,
xla_cpu_use_acl: builtins.bool | None = ...,
xla_cpu_strict_dot_conv_math: builtins.bool | None = ...,
xla_gpu_use_runtime_fusion: builtins.bool | None = ...,
@@ -1020,6 +1133,7 @@ class DebugOptions(google.protobuf.message.Message):
xla_gpu_exhaustive_tiling_search: builtins.bool | None = ...,
xla_gpu_enable_triton_softmax_fusion: builtins.bool | None = ...,
xla_gpu_enable_priority_fusion: builtins.bool | None = ...,
xla_gpu_enable_triton_softmax_priority_fusion: builtins.bool | None = ...,
xla_gpu_dump_autotune_results_to: builtins.str | None = ...,
xla_gpu_load_autotune_results_from: builtins.str | None = ...,
xla_gpu_target_config_filename: builtins.str | None = ...,
@@ -1027,7 +1141,8 @@ class DebugOptions(google.protobuf.message.Message):
xla_gpu_auto_spmd_partitioning_memory_budget_ratio: builtins.float | None = ...,
xla_gpu_triton_gemm_disable_reduced_precision_reduction: builtins.bool | None = ...,
xla_gpu_triton_fusion_level: builtins.int | None = ...,
xla_gpu_dump_autotuned_triton_fusions: builtins.bool | None = ...,
xla_gpu_dump_autotuned_gemm_fusions: builtins.bool | None = ...,
xla_gpu_override_gemm_autotuner: builtins.str | None = ...,
xla_gpu_copy_insertion_use_region_analysis: builtins.bool | None = ...,
xla_gpu_collect_cost_model_stats: builtins.bool | None = ...,
xla_gpu_enable_split_k_autotuning: builtins.bool | None = ...,
@@ -1036,6 +1151,7 @@ class DebugOptions(google.protobuf.message.Message):
xla_gpu_mock_custom_calls: builtins.bool | None = ...,
xla_gpu_cublas_fallback: builtins.bool | None = ...,
xla_gpu_enable_while_loop_double_buffering: builtins.bool | None = ...,
xla_gpu_enable_while_loop_unrolling: global___DebugOptions.WhileLoopUnrolling.ValueType | None = ...,
xla_gpu_ensure_minor_dot_contraction_dims: builtins.bool | None = ...,
xla_gpu_filter_kernels_spilling_registers_on_autotuning: builtins.bool | None = ...,
xla_debug_buffer_assignment_show_max: builtins.int | None = ...,
@@ -1044,11 +1160,32 @@ class DebugOptions(google.protobuf.message.Message):
xla_gpu_threshold_for_windowed_einsum_mib: builtins.int | None = ...,
xla_gpu_enable_triton_hopper: builtins.bool | None = ...,
xla_gpu_enable_nccl_user_buffers: builtins.bool | None = ...,
xla_gpu_enable_nccl_comm_splitting: builtins.bool | None = ...,
xla_gpu_enable_nccl_per_stream_comms: builtins.bool | None = ...,
xla_gpu_enable_libnvptxcompiler: builtins.bool | None = ...,
xla_gpu_enable_dot_strength_reduction: builtins.bool | None = ...,
xla_gpu_multi_streamed_windowed_einsum: builtins.bool | None = ...,
xla_gpu_enable_bf16_6way_gemm: builtins.bool | None = ...,
xla_gpu_enable_bf16_3way_gemm: builtins.bool | None = ...,
xla_gpu_nccl_collective_max_nchannels: builtins.int | None = ...,
xla_gpu_nccl_p2p_max_nchannels: builtins.int | None = ...,
xla_gpu_enable_mlir_emitters: builtins.bool | None = ...,
xla_gpu_max_mlir_kernels: builtins.int | None = ...,
xla_gpu_skip_mlir_kernels: builtins.int | None = ...,
xla_gpu_gemm_rewrite_size_threshold: builtins.int | None = ...,
xla_gpu_require_complete_aot_autotune_results: builtins.bool | None = ...,
xla_gpu_cudnn_gemm_fusion_level: builtins.int | None = ...,
xla_gpu_use_memcpy_local_p2p: builtins.bool | None = ...,
xla_gpu_autotune_max_solutions: builtins.int | None = ...,
xla_dump_large_constants: builtins.bool | None = ...,
xla_gpu_verify_triton_fusion_numerics: builtins.bool | None = ...,
xla_gpu_dump_autotune_logs_to: builtins.str | None = ...,
xla_reduce_window_rewrite_base_length: builtins.int | None = ...,
xla_gpu_enable_host_memory_offloading: builtins.bool | None = ...,
xla_gpu_exclude_nondeterministic_ops: builtins.bool | None = ...,
xla_backend_extra_options: collections.abc.Mapping[builtins.str, builtins.str] | None = ...,
) -> None: ...
def ClearField(self, field_name: typing.Literal["xla_allow_excess_precision", b"xla_allow_excess_precision", "xla_allow_scalar_index_dynamic_ops", b"xla_allow_scalar_index_dynamic_ops", "xla_backend_extra_options", b"xla_backend_extra_options", "xla_backend_optimization_level", b"xla_backend_optimization_level", "xla_cpu_enable_custom_matmul_tiling", b"xla_cpu_enable_custom_matmul_tiling", "xla_cpu_enable_experimental_deallocation", b"xla_cpu_enable_experimental_deallocation", "xla_cpu_enable_fast_math", b"xla_cpu_enable_fast_math", "xla_cpu_enable_fast_min_max", b"xla_cpu_enable_fast_min_max", "xla_cpu_enable_mlir_fusion_outlining", b"xla_cpu_enable_mlir_fusion_outlining", "xla_cpu_enable_mlir_tiling_and_fusion", b"xla_cpu_enable_mlir_tiling_and_fusion", "xla_cpu_enable_xprof_traceme", b"xla_cpu_enable_xprof_traceme", "xla_cpu_fast_math_honor_division", b"xla_cpu_fast_math_honor_division", "xla_cpu_fast_math_honor_functions", b"xla_cpu_fast_math_honor_functions", "xla_cpu_fast_math_honor_infs", b"xla_cpu_fast_math_honor_infs", "xla_cpu_fast_math_honor_nans", b"xla_cpu_fast_math_honor_nans", "xla_cpu_matmul_tiling_k_dim", b"xla_cpu_matmul_tiling_k_dim", "xla_cpu_matmul_tiling_m_dim", b"xla_cpu_matmul_tiling_m_dim", "xla_cpu_matmul_tiling_n_dim", b"xla_cpu_matmul_tiling_n_dim", "xla_cpu_multi_thread_eigen", b"xla_cpu_multi_thread_eigen", "xla_cpu_sparse_cuda_threads", b"xla_cpu_sparse_cuda_threads", "xla_cpu_strict_dot_conv_math", b"xla_cpu_strict_dot_conv_math", "xla_cpu_use_acl", b"xla_cpu_use_acl", "xla_cpu_use_mkl_dnn", b"xla_cpu_use_mkl_dnn", "xla_cpu_use_xla_runtime", b"xla_cpu_use_xla_runtime", "xla_debug_buffer_assignment_show_max", b"xla_debug_buffer_assignment_show_max", "xla_detailed_logging", b"xla_detailed_logging", "xla_disable_all_hlo_passes", b"xla_disable_all_hlo_passes", "xla_disable_hlo_passes", b"xla_disable_hlo_passes", "xla_dump_compress_protos", b"xla_dump_compress_protos", "xla_dump_disable_metadata", b"xla_dump_disable_metadata", "xla_dump_enable_mlir_pretty_form", b"xla_dump_enable_mlir_pretty_form", "xla_dump_fusion_visualization", b"xla_dump_fusion_visualization", "xla_dump_hlo_as_dot", b"xla_dump_hlo_as_dot", "xla_dump_hlo_as_html", b"xla_dump_hlo_as_html", "xla_dump_hlo_as_long_text", b"xla_dump_hlo_as_long_text", "xla_dump_hlo_as_proto", b"xla_dump_hlo_as_proto", "xla_dump_hlo_as_text", b"xla_dump_hlo_as_text", "xla_dump_hlo_as_url", b"xla_dump_hlo_as_url", "xla_dump_hlo_module_re", b"xla_dump_hlo_module_re", "xla_dump_hlo_pass_re", b"xla_dump_hlo_pass_re", "xla_dump_hlo_pipeline_re", b"xla_dump_hlo_pipeline_re", "xla_dump_hlo_snapshots", b"xla_dump_hlo_snapshots", "xla_dump_include_timestamp", b"xla_dump_include_timestamp", "xla_dump_latency_hiding_schedule", b"xla_dump_latency_hiding_schedule", "xla_dump_max_hlo_modules", b"xla_dump_max_hlo_modules", "xla_dump_module_metadata", b"xla_dump_module_metadata", "xla_dump_to", b"xla_dump_to", "xla_eliminate_hlo_implicit_broadcast", b"xla_eliminate_hlo_implicit_broadcast", "xla_embed_ir_in_executable", b"xla_embed_ir_in_executable", "xla_enable_dumping", b"xla_enable_dumping", "xla_enable_hlo_passes_only", b"xla_enable_hlo_passes_only", "xla_force_host_platform_device_count", b"xla_force_host_platform_device_count", "xla_gpu_algorithm_denylist_path", b"xla_gpu_algorithm_denylist_path", "xla_gpu_all_gather_combine_threshold_bytes", b"xla_gpu_all_gather_combine_threshold_bytes", "xla_gpu_all_reduce_blueconnect_num_devices_per_host", b"xla_gpu_all_reduce_blueconnect_num_devices_per_host", "xla_gpu_all_reduce_combine_threshold_bytes", b"xla_gpu_all_reduce_combine_threshold_bytes", "xla_gpu_all_reduce_contiguous", b"xla_gpu_all_reduce_contiguous", "xla_gpu_asm_extra_flags", b"xla_gpu_asm_extra_flags", "xla_gpu_auto_spmd_partitioning_memory_budget_gb", b"xla_gpu_auto_spmd_partitioning_memory_budget_gb", "xla_gpu_auto_spmd_partitioning_memory_budget_ratio", b"xla_gpu_auto_spmd_partitioning_memory_budget_ratio", "xla_gpu_autotune_level", b"xla_gpu_autotune_level", "xla_gpu_collect_cost_model_stats", b"xla_gpu_collect_cost_model_stats", "xla_gpu_collective_inflation_factor", b"xla_gpu_collective_inflation_factor", "xla_gpu_collective_permute_decomposer_threshold", b"xla_gpu_collective_permute_decomposer_threshold", "xla_gpu_copy_insertion_use_region_analysis", b"xla_gpu_copy_insertion_use_region_analysis", "xla_gpu_crash_on_verification_failures", b"xla_gpu_crash_on_verification_failures", "xla_gpu_cublas_fallback", b"xla_gpu_cublas_fallback", "xla_gpu_cuda_data_dir", b"xla_gpu_cuda_data_dir", "xla_gpu_deterministic_ops", b"xla_gpu_deterministic_ops", "xla_gpu_disable_gpuasm_optimizations", b"xla_gpu_disable_gpuasm_optimizations", "xla_gpu_dump_autotune_results_to", b"xla_gpu_dump_autotune_results_to", "xla_gpu_dump_autotuned_triton_fusions", b"xla_gpu_dump_autotuned_triton_fusions", "xla_gpu_dump_llvmir", b"xla_gpu_dump_llvmir", "xla_gpu_enable_address_computation_fusion", b"xla_gpu_enable_address_computation_fusion", "xla_gpu_enable_aLine truncated
def ClearField(self, field_name: typing.Literal["xla_allow_excess_precision", b"xla_allow_excess_precision", "xla_allow_scalar_index_dynamic_ops", b"xla_allow_scalar_index_dynamic_ops", "xla_backend_extra_options", b"xla_backend_extra_options", "xla_backend_optimization_level", b"xla_backend_optimization_level", "xla_cpu_enable_custom_matmul_tiling", b"xla_cpu_enable_custom_matmul_tiling", "xla_cpu_enable_experimental_deallocation", b"xla_cpu_enable_experimental_deallocation", "xla_cpu_enable_fast_math", b"xla_cpu_enable_fast_math", "xla_cpu_enable_fast_min_max", b"xla_cpu_enable_fast_min_max", "xla_cpu_enable_mlir_fusion_outlining", b"xla_cpu_enable_mlir_fusion_outlining", "xla_cpu_enable_mlir_tiling_and_fusion", b"xla_cpu_enable_mlir_tiling_and_fusion", "xla_cpu_enable_xprof_traceme", b"xla_cpu_enable_xprof_traceme", "xla_cpu_fast_math_honor_division", b"xla_cpu_fast_math_honor_division", "xla_cpu_fast_math_honor_functions", b"xla_cpu_fast_math_honor_functions", "xla_cpu_fast_math_honor_infs", b"xla_cpu_fast_math_honor_infs", "xla_cpu_fast_math_honor_nans", b"xla_cpu_fast_math_honor_nans", "xla_cpu_matmul_tiling_k_dim", b"xla_cpu_matmul_tiling_k_dim", "xla_cpu_matmul_tiling_m_dim", b"xla_cpu_matmul_tiling_m_dim", "xla_cpu_matmul_tiling_n_dim", b"xla_cpu_matmul_tiling_n_dim", "xla_cpu_multi_thread_eigen", b"xla_cpu_multi_thread_eigen", "xla_cpu_strict_dot_conv_math", b"xla_cpu_strict_dot_conv_math", "xla_cpu_use_acl", b"xla_cpu_use_acl", "xla_cpu_use_mkl_dnn", b"xla_cpu_use_mkl_dnn", "xla_cpu_use_thunk_runtime", b"xla_cpu_use_thunk_runtime", "xla_debug_buffer_assignment_show_max", b"xla_debug_buffer_assignment_show_max", "xla_detailed_logging", b"xla_detailed_logging", "xla_disable_all_hlo_passes", b"xla_disable_all_hlo_passes", "xla_disable_hlo_passes", b"xla_disable_hlo_passes", "xla_dump_compress_protos", b"xla_dump_compress_protos", "xla_dump_disable_metadata", b"xla_dump_disable_metadata", "xla_dump_enable_mlir_pretty_form", b"xla_dump_enable_mlir_pretty_form", "xla_dump_fusion_visualization", b"xla_dump_fusion_visualization", "xla_dump_hlo_as_dot", b"xla_dump_hlo_as_dot", "xla_dump_hlo_as_html", b"xla_dump_hlo_as_html", "xla_dump_hlo_as_long_text", b"xla_dump_hlo_as_long_text", "xla_dump_hlo_as_proto", b"xla_dump_hlo_as_proto", "xla_dump_hlo_as_text", b"xla_dump_hlo_as_text", "xla_dump_hlo_as_url", b"xla_dump_hlo_as_url", "xla_dump_hlo_module_re", b"xla_dump_hlo_module_re", "xla_dump_hlo_pass_re", b"xla_dump_hlo_pass_re", "xla_dump_hlo_pipeline_re", b"xla_dump_hlo_pipeline_re", "xla_dump_hlo_snapshots", b"xla_dump_hlo_snapshots", "xla_dump_include_timestamp", b"xla_dump_include_timestamp", "xla_dump_large_constants", b"xla_dump_large_constants", "xla_dump_latency_hiding_schedule", b"xla_dump_latency_hiding_schedule", "xla_dump_max_hlo_modules", b"xla_dump_max_hlo_modules", "xla_dump_module_metadata", b"xla_dump_module_metadata", "xla_dump_to", b"xla_dump_to", "xla_eliminate_hlo_implicit_broadcast", b"xla_eliminate_hlo_implicit_broadcast", "xla_embed_ir_in_executable", b"xla_embed_ir_in_executable", "xla_enable_dumping", b"xla_enable_dumping", "xla_enable_hlo_passes_only", b"xla_enable_hlo_passes_only", "xla_force_host_platform_device_count", b"xla_force_host_platform_device_count", "xla_gpu_algorithm_denylist_path", b"xla_gpu_algorithm_denylist_path", "xla_gpu_all_gather_combine_threshold_bytes", b"xla_gpu_all_gather_combine_threshold_bytes", "xla_gpu_all_reduce_blueconnect_num_devices_per_host", b"xla_gpu_all_reduce_blueconnect_num_devices_per_host", "xla_gpu_all_reduce_combine_threshold_bytes", b"xla_gpu_all_reduce_combine_threshold_bytes", "xla_gpu_all_reduce_contiguous", b"xla_gpu_all_reduce_contiguous", "xla_gpu_asm_extra_flags", b"xla_gpu_asm_extra_flags", "xla_gpu_auto_spmd_partitioning_memory_budget_gb", b"xla_gpu_auto_spmd_partitioning_memory_budget_gb", "xla_gpu_auto_spmd_partitioning_memory_budget_ratio", b"xla_gpu_auto_spmd_partitioning_memory_budget_ratio", "xla_gpu_autotune_level", b"xla_gpu_autotune_level", "xla_gpu_autotune_max_solutions", b"xla_gpu_autotune_max_solutions", "xla_gpu_collect_cost_model_stats", b"xla_gpu_collect_cost_model_stats", "xla_gpu_collective_inflation_factor", b"xla_gpu_collective_inflation_factor", "xla_gpu_collective_permute_decomposer_threshold", b"xla_gpu_collective_permute_decomposer_threshold", "xla_gpu_copy_insertion_use_region_analysis", b"xla_gpu_copy_insertion_use_region_analysis", "xla_gpu_crash_on_verification_failures", b"xla_gpu_crash_on_verification_failures", "xla_gpu_cublas_fallback", b"xla_gpu_cublas_fallback", "xla_gpu_cuda_data_dir", b"xla_gpu_cuda_data_dir", "xla_gpu_cudnn_gemm_fusion_level", b"xla_gpu_cudnn_gemm_fusion_level", "xla_gpu_deterministic_ops", b"xla_gpu_deterministic_ops", "xla_gpu_disable_async_collectives", b"xla_gpu_disable_async_collectives", "xla_gpu_disable_gpuasm_optimizations", b"xla_gpu_disable_gpuasm_optimizations", "xla_gpu_dump_autotune_logs_to", b"xla_gpu_dump_autotune_logs_to", "xla_gpu_dump_autotune_results_Line truncated
global___DebugOptions = DebugOptions
@@ -1104,7 +1241,7 @@ class ExecutionOptions(google.protobuf.message.Message):
will have an effect on every platform.
When adding new fields, keep in mind that boolean fields default to false.
Next id: 23.
Next id: 24.
"""
DESCRIPTOR: google.protobuf.descriptor.Descriptor
@@ -1123,6 +1260,7 @@ class ExecutionOptions(google.protobuf.message.Message):
AUTO_SPMD_PARTITIONING_MESH_SHAPE_FIELD_NUMBER: builtins.int
AUTO_SPMD_PARTITIONING_MESH_IDS_FIELD_NUMBER: builtins.int
DEDUPLICATE_HLO_FIELD_NUMBER: builtins.int
ALLOW_SPMD_SHARDING_PROPAGATION_TO_PARAMETERS_FIELD_NUMBER: builtins.int
ALLOW_SPMD_SHARDING_PROPAGATION_TO_OUTPUT_FIELD_NUMBER: builtins.int
PARAM_REQUIRES_BROADCAST_VIA_COLLECTIVES_FIELD_NUMBER: builtins.int
ALLOW_SEPARATE_SHARDING_PROGRAMS_FIELD_NUMBER: builtins.int
@@ -1207,6 +1345,20 @@ class ExecutionOptions(google.protobuf.message.Message):
use_auto_spmd_partitioning=true.
"""
@property
def allow_spmd_sharding_propagation_to_parameters(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.bool]:
"""Allows sharding propagation to propagate to the parameters. This changes
the input shape of the computation (which is undesirable), but it can be
used to allow to run partial compilation to determine what would be the
input sharding of a computation if XLA would be allowed to propagate the
sharding which can be used by higher level framework as a way to query
intermediate sharding of operations when multiple computation would be
chained and merged together.
This is a vector of bool, because the user can control which parameters can
have the sharding substituted. If only one boolean value is passed in the
vector that is interpreted as the value to be applied for every parameter.
"""
@property
def allow_spmd_sharding_propagation_to_output(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.bool]:
"""Allows sharding propagation to propagate to the outputs. This changes the
@@ -1249,6 +1401,7 @@ class ExecutionOptions(google.protobuf.message.Message):
auto_spmd_partitioning_mesh_shape: collections.abc.Iterable[builtins.int] | None = ...,
auto_spmd_partitioning_mesh_ids: collections.abc.Iterable[builtins.int] | None = ...,
deduplicate_hlo: builtins.bool | None = ...,
allow_spmd_sharding_propagation_to_parameters: collections.abc.Iterable[builtins.bool] | None = ...,
allow_spmd_sharding_propagation_to_output: collections.abc.Iterable[builtins.bool] | None = ...,
param_requires_broadcast_via_collectives: collections.abc.Iterable[builtins.bool] | None = ...,
allow_separate_sharding_programs: builtins.bool | None = ...,
@@ -1257,7 +1410,7 @@ class ExecutionOptions(google.protobuf.message.Message):
device_memory_size: builtins.int | None = ...,
) -> None: ...
def HasField(self, field_name: typing.Literal["debug_options", b"debug_options", "device_assignment", b"device_assignment", "shape_with_output_layout", b"shape_with_output_layout"]) -> builtins.bool: ...
def ClearField(self, field_name: typing.Literal["alias_passthrough_params", b"alias_passthrough_params", "allow_separate_sharding_programs", b"allow_separate_sharding_programs", "allow_spmd_sharding_propagation_to_output", b"allow_spmd_sharding_propagation_to_output", "auto_spmd_partitioning_mesh_ids", b"auto_spmd_partitioning_mesh_ids", "auto_spmd_partitioning_mesh_shape", b"auto_spmd_partitioning_mesh_shape", "debug_options", b"debug_options", "deduplicate_hlo", b"deduplicate_hlo", "device_assignment", b"device_assignment", "device_handles", b"device_handles", "device_memory_size", b"device_memory_size", "fdo_profile", b"fdo_profile", "launch_id", b"launch_id", "num_partitions", b"num_partitions", "num_replicas", b"num_replicas", "param_requires_broadcast_via_collectives", b"param_requires_broadcast_via_collectives", "seed", b"seed", "shape_with_output_layout", b"shape_with_output_layout", "shardable_value_update_pairs", b"shardable_value_update_pairs", "use_auto_spmd_partitioning", b"use_auto_spmd_partitioning", "use_spmd_partitioning", b"use_spmd_partitioning"]) -> None: ...
def ClearField(self, field_name: typing.Literal["alias_passthrough_params", b"alias_passthrough_params", "allow_separate_sharding_programs", b"allow_separate_sharding_programs", "allow_spmd_sharding_propagation_to_output", b"allow_spmd_sharding_propagation_to_output", "allow_spmd_sharding_propagation_to_parameters", b"allow_spmd_sharding_propagation_to_parameters", "auto_spmd_partitioning_mesh_ids", b"auto_spmd_partitioning_mesh_ids", "auto_spmd_partitioning_mesh_shape", b"auto_spmd_partitioning_mesh_shape", "debug_options", b"debug_options", "deduplicate_hlo", b"deduplicate_hlo", "device_assignment", b"device_assignment", "device_handles", b"device_handles", "device_memory_size", b"device_memory_size", "fdo_profile", b"fdo_profile", "launch_id", b"launch_id", "num_partitions", b"num_partitions", "num_replicas", b"num_replicas", "param_requires_broadcast_via_collectives", b"param_requires_broadcast_via_collectives", "seed", b"seed", "shape_with_output_layout", b"shape_with_output_layout", "shardable_value_update_pairs", b"shardable_value_update_pairs", "use_auto_spmd_partitioning", b"use_auto_spmd_partitioning", "use_spmd_partitioning", b"use_spmd_partitioning"]) -> None: ...
global___ExecutionOptions = ExecutionOptions
@@ -1266,7 +1419,7 @@ class HloModuleConfigProto(google.protobuf.message.Message):
"""Serialization of HloModuleConfig. See the C++ class definition for
descriptions of each field.
There are no guarantees of backwards or forwards compatibility.
Next id: 33.
Next id: 34.
"""
DESCRIPTOR: google.protobuf.descriptor.Descriptor
@@ -1394,6 +1547,7 @@ class HloModuleConfigProto(google.protobuf.message.Message):
MEMORY_SPACE_ASSIGNMENT_CONFIG_FIELD_NUMBER: builtins.int
PHASE_ORDERING_CONFIG_FIELD_NUMBER: builtins.int
PHASE_INDEX_FIELD_NUMBER: builtins.int
ALLOW_SPMD_SHARDING_PROPAGATION_TO_PARAMETERS_FIELD_NUMBER: builtins.int
ALLOW_SPMD_SHARDING_PROPAGATION_TO_OUTPUT_FIELD_NUMBER: builtins.int
ANALYSIS_ALLOWANCE_MAP_FIELD_NUMBER: builtins.int
MATRIX_UNIT_OPERAND_PRECISION_FIELD_NUMBER: builtins.int
@@ -1441,6 +1595,8 @@ class HloModuleConfigProto(google.protobuf.message.Message):
@property
def phase_ordering_config(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___HloModuleConfigProto.BoolList]: ...
@property
def allow_spmd_sharding_propagation_to_parameters(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.bool]: ...
@property
def allow_spmd_sharding_propagation_to_output(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.bool]: ...
@property
def analysis_allowance_map(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.int]: ...
@@ -1473,6 +1629,7 @@ class HloModuleConfigProto(google.protobuf.message.Message):
memory_space_assignment_config: collections.abc.Iterable[builtins.int] | None = ...,
phase_ordering_config: collections.abc.Iterable[global___HloModuleConfigProto.BoolList] | None = ...,
phase_index: builtins.int | None = ...,
allow_spmd_sharding_propagation_to_parameters: collections.abc.Iterable[builtins.bool] | None = ...,
allow_spmd_sharding_propagation_to_output: collections.abc.Iterable[builtins.bool] | None = ...,
analysis_allowance_map: collections.abc.Mapping[builtins.str, builtins.int] | None = ...,
matrix_unit_operand_precision: tensorflow.compiler.xla.xla_data_pb2.PrecisionConfig.Precision.ValueType | None = ...,
@@ -1480,7 +1637,7 @@ class HloModuleConfigProto(google.protobuf.message.Message):
device_memory_size: builtins.int | None = ...,
) -> None: ...
def HasField(self, field_name: typing.Literal["debug_options", b"debug_options", "entry_computation_layout", b"entry_computation_layout", "static_device_assignment", b"static_device_assignment"]) -> builtins.bool: ...
def ClearField(self, field_name: typing.Literal["alias_passthrough_params", b"alias_passthrough_params", "allow_separate_sharding_programs", b"allow_separate_sharding_programs", "allow_spmd_sharding_propagation_to_output", b"allow_spmd_sharding_propagation_to_output", "analysis_allowance_map", b"analysis_allowance_map", "auto_spmd_partitioning_mesh_ids", b"auto_spmd_partitioning_mesh_ids", "auto_spmd_partitioning_mesh_shape", b"auto_spmd_partitioning_mesh_shape", "content_aware_computation_sorting", b"content_aware_computation_sorting", "debug_options", b"debug_options", "deduplicate_hlo", b"deduplicate_hlo", "device_memory_size", b"device_memory_size", "device_type", b"device_type", "dot_config", b"dot_config", "entry_computation_layout", b"entry_computation_layout", "fdo_profile", b"fdo_profile", "fusion_config", b"fusion_config", "fusion_config_collection", b"fusion_config_collection", "intra_op_parallelism_threads", b"intra_op_parallelism_threads", "launch_id", b"launch_id", "layout_config", b"layout_config", "matrix_unit_operand_precision", b"matrix_unit_operand_precision", "memory_space_assignment_config", b"memory_space_assignment_config", "num_partitions", b"num_partitions", "param_requires_broadcast_via_collectives", b"param_requires_broadcast_via_collectives", "phase_index", b"phase_index", "phase_ordering_config", b"phase_ordering_config", "replica_count", b"replica_count", "seed", b"seed", "shardable_value_update_pairs", b"shardable_value_update_pairs", "static_device_assignment", b"static_device_assignment", "use_auto_spmd_partitioning", b"use_auto_spmd_partitioning", "use_spmd_partitioning", b"use_spmd_partitioning"]) -> None: ...
def ClearField(self, field_name: typing.Literal["alias_passthrough_params", b"alias_passthrough_params", "allow_separate_sharding_programs", b"allow_separate_sharding_programs", "allow_spmd_sharding_propagation_to_output", b"allow_spmd_sharding_propagation_to_output", "allow_spmd_sharding_propagation_to_parameters", b"allow_spmd_sharding_propagation_to_parameters", "analysis_allowance_map", b"analysis_allowance_map", "auto_spmd_partitioning_mesh_ids", b"auto_spmd_partitioning_mesh_ids", "auto_spmd_partitioning_mesh_shape", b"auto_spmd_partitioning_mesh_shape", "content_aware_computation_sorting", b"content_aware_computation_sorting", "debug_options", b"debug_options", "deduplicate_hlo", b"deduplicate_hlo", "device_memory_size", b"device_memory_size", "device_type", b"device_type", "dot_config", b"dot_config", "entry_computation_layout", b"entry_computation_layout", "fdo_profile", b"fdo_profile", "fusion_config", b"fusion_config", "fusion_config_collection", b"fusion_config_collection", "intra_op_parallelism_threads", b"intra_op_parallelism_threads", "launch_id", b"launch_id", "layout_config", b"layout_config", "matrix_unit_operand_precision", b"matrix_unit_operand_precision", "memory_space_assignment_config", b"memory_space_assignment_config", "num_partitions", b"num_partitions", "param_requires_broadcast_via_collectives", b"param_requires_broadcast_via_collectives", "phase_index", b"phase_index", "phase_ordering_config", b"phase_ordering_config", "replica_count", b"replica_count", "seed", b"seed", "shardable_value_update_pairs", b"shardable_value_update_pairs", "static_device_assignment", b"static_device_assignment", "use_auto_spmd_partitioning", b"use_auto_spmd_partitioning", "use_spmd_partitioning", b"use_spmd_partitioning"]) -> None: ...
global___HloModuleConfigProto = HloModuleConfigProto
stubs/tensorflow/tensorflow/core/framework/dataset_options_pb2.pyi
+8 -3
View File
@@ -93,7 +93,7 @@ global___ExternalStatePolicy = ExternalStatePolicy
@typing.final
class AutotuneOptions(google.protobuf.message.Message):
"""next: 5"""
"""next: 6"""
DESCRIPTOR: google.protobuf.descriptor.Descriptor
@@ -101,10 +101,12 @@ class AutotuneOptions(google.protobuf.message.Message):
CPU_BUDGET_FIELD_NUMBER: builtins.int
RAM_BUDGET_FIELD_NUMBER: builtins.int
AUTOTUNE_ALGORITHM_FIELD_NUMBER: builtins.int
INITIAL_PARALLELISM_FIELD_NUMBER: builtins.int
enabled: builtins.bool
cpu_budget: builtins.int
ram_budget: builtins.int
autotune_algorithm: tensorflow.core.framework.model_pb2.AutotuneAlgorithm.ValueType
initial_parallelism: builtins.int
def __init__(
self,
*,
@@ -112,9 +114,10 @@ class AutotuneOptions(google.protobuf.message.Message):
cpu_budget: builtins.int | None = ...,
ram_budget: builtins.int | None = ...,
autotune_algorithm: tensorflow.core.framework.model_pb2.AutotuneAlgorithm.ValueType | None = ...,
initial_parallelism: builtins.int | None = ...,
) -> None: ...
def HasField(self, field_name: typing.Literal["autotune_algorithm", b"autotune_algorithm", "cpu_budget", b"cpu_budget", "enabled", b"enabled", "optional_autotune_algorithm", b"optional_autotune_algorithm", "optional_cpu_budget", b"optional_cpu_budget", "optional_enabled", b"optional_enabled", "optional_ram_budget", b"optional_ram_budget", "ram_budget", b"ram_budget"]) -> builtins.bool: ...
def ClearField(self, field_name: typing.Literal["autotune_algorithm", b"autotune_algorithm", "cpu_budget", b"cpu_budget", "enabled", b"enabled", "optional_autotune_algorithm", b"optional_autotune_algorithm", "optional_cpu_budget", b"optional_cpu_budget", "optional_enabled", b"optional_enabled", "optional_ram_budget", b"optional_ram_budget", "ram_budget", b"ram_budget"]) -> None: ...
def HasField(self, field_name: typing.Literal["autotune_algorithm", b"autotune_algorithm", "cpu_budget", b"cpu_budget", "enabled", b"enabled", "initial_parallelism", b"initial_parallelism", "optional_autotune_algorithm", b"optional_autotune_algorithm", "optional_cpu_budget", b"optional_cpu_budget", "optional_enabled", b"optional_enabled", "optional_initial_parallelism", b"optional_initial_parallelism", "optional_ram_budget", b"optional_ram_budget", "ram_budget", b"ram_budget"]) -> builtins.bool: ...
def ClearField(self, field_name: typing.Literal["autotune_algorithm", b"autotune_algorithm", "cpu_budget", b"cpu_budget", "enabled", b"enabled", "initial_parallelism", b"initial_parallelism", "optional_autotune_algorithm", b"optional_autotune_algorithm", "optional_cpu_budget", b"optional_cpu_budget", "optional_enabled", b"optional_enabled", "optional_initial_parallelism", b"optional_initial_parallelism", "optional_ram_budget", b"optional_ram_budget", "ram_budget", b"ram_budget"]) -> None: ...
@typing.overload
def WhichOneof(self, oneof_group: typing.Literal["optional_autotune_algorithm", b"optional_autotune_algorithm"]) -> typing.Literal["autotune_algorithm"] | None: ...
@typing.overload
@@ -122,6 +125,8 @@ class AutotuneOptions(google.protobuf.message.Message):
@typing.overload
def WhichOneof(self, oneof_group: typing.Literal["optional_enabled", b"optional_enabled"]) -> typing.Literal["enabled"] | None: ...
@typing.overload
def WhichOneof(self, oneof_group: typing.Literal["optional_initial_parallelism", b"optional_initial_parallelism"]) -> typing.Literal["initial_parallelism"] | None: ...
@typing.overload
def WhichOneof(self, oneof_group: typing.Literal["optional_ram_budget", b"optional_ram_budget"]) -> typing.Literal["ram_budget"] | None: ...
global___AutotuneOptions = AutotuneOptions
stubs/tensorflow/tensorflow/core/framework/resource_handle_pb2.pyi
+4 -1
View File
@@ -33,8 +33,11 @@ class ResourceHandleProto(google.protobuf.message.Message):
DTYPE_FIELD_NUMBER: builtins.int
SHAPE_FIELD_NUMBER: builtins.int
dtype: tensorflow.core.framework.types_pb2.DataType.ValueType
"""Data type of the tensor."""
@property
def shape(self) -> tensorflow.core.framework.tensor_shape_pb2.TensorShapeProto: ...
def shape(self) -> tensorflow.core.framework.tensor_shape_pb2.TensorShapeProto:
"""Shape of the tensor."""
def __init__(
self,
*,
stubs/tensorflow/tensorflow/core/framework/tensor_pb2.pyi
+1
View File
@@ -41,6 +41,7 @@ class TensorProto(google.protobuf.message.Message):
UINT64_VAL_FIELD_NUMBER: builtins.int
FLOAT8_VAL_FIELD_NUMBER: builtins.int
dtype: tensorflow.core.framework.types_pb2.DataType.ValueType
"""Data type of the tensor."""
version_number: builtins.int
"""Only one of the representations below is set, one of "tensor_contents" and
the "xxx_val" attributes. We are not using oneof because as oneofs cannot
stubs/tensorflow/tensorflow/core/protobuf/config_pb2.pyi
+61 -1
View File
@@ -90,6 +90,49 @@ class GPUOptions(google.protobuf.message.Message):
) -> None: ...
def ClearField(self, field_name: typing.Literal["device_ordinal", b"device_ordinal", "memory_limit_mb", b"memory_limit_mb", "priority", b"priority"]) -> None: ...
@typing.final
class StreamMergeOptions(google.protobuf.message.Message):
"""Whether to merge data transfer streams into the compute stream in the
same stream group. Stream merging helps reduce the overhead caused by
stream synchronization, especially when data transfers are frequent. For
example, setting "merge_host_to_device_stream = true" will make the
compute stream responsible for both computation and host to device memory
copy.
"""
DESCRIPTOR: google.protobuf.descriptor.Descriptor
MERGE_HOST_TO_DEVICE_STREAM_FIELD_NUMBER: builtins.int
MERGE_DEVICE_TO_HOST_STREAM_FIELD_NUMBER: builtins.int
MERGE_DEVICE_TO_DEVICE_STREAM_FIELD_NUMBER: builtins.int
merge_host_to_device_stream: builtins.bool
"""If true, the compute stream will be used for host_to_device copy as
well. It's no longer necessary to record an event before the copy to
let the copy stream wait for the compute stream to finish. There is
also no need to wait for the copy to complete before executing the
callback function.
"""
merge_device_to_host_stream: builtins.bool
"""If true, the compute stream will be used for device_to_host copy as
well. It's no longer necessary to record an event before the copy to
let the copy stream wait for the compute stream to finish.
"""
merge_device_to_device_stream: builtins.bool
"""If true, the compute stream will be used for device_to_device copy as
well. It's no longer necessary to record an event before the copy to
let the copy stream wait for the compute stream of the sending device
to finish. There is also no need to wait for the compute stream of the
receiving device to finish if the copy is within the same device.
"""
def __init__(
self,
*,
merge_host_to_device_stream: builtins.bool | None = ...,
merge_device_to_host_stream: builtins.bool | None = ...,
merge_device_to_device_stream: builtins.bool | None = ...,
) -> None: ...
def ClearField(self, field_name: typing.Literal["merge_device_to_device_stream", b"merge_device_to_device_stream", "merge_device_to_host_stream", b"merge_device_to_host_stream", "merge_host_to_device_stream", b"merge_host_to_device_stream"]) -> None: ...
VIRTUAL_DEVICES_FIELD_NUMBER: builtins.int
NUM_VIRTUAL_DEVICES_PER_GPU_FIELD_NUMBER: builtins.int
USE_UNIFIED_MEMORY_FIELD_NUMBER: builtins.int
@@ -105,6 +148,9 @@ class GPUOptions(google.protobuf.message.Message):
GPU_HOST_MEM_LIMIT_IN_MB_FIELD_NUMBER: builtins.int
GPU_HOST_MEM_DISALLOW_GROWTH_FIELD_NUMBER: builtins.int
GPU_SYSTEM_MEMORY_SIZE_IN_MB_FIELD_NUMBER: builtins.int
POPULATE_PJRT_GPU_CLIENT_CREATION_INFO_FIELD_NUMBER: builtins.int
NODE_ID_FIELD_NUMBER: builtins.int
STREAM_MERGE_OPTIONS_FIELD_NUMBER: builtins.int
num_virtual_devices_per_gpu: builtins.int
"""The number of virtual devices to create on each visible GPU. The
available memory will be split equally among all virtual devices. If the
@@ -200,6 +246,14 @@ class GPUOptions(google.protobuf.message.Message):
system memory size for better resource estimation of multi-tenancy(one
gpu with multiple model) use case.
"""
populate_pjrt_gpu_client_creation_info: builtins.bool
"""If true, save information needed for created a PjRt GPU client for
creating a client with remote devices.
"""
node_id: builtins.int
"""node_id for use when creating a PjRt GPU client with remote devices,
which enumerates jobs*tasks from a ServerDef.
"""
@property
def virtual_devices(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___GPUOptions.Experimental.VirtualDevices]:
"""The multi virtual device settings. If empty (not set), it will create
@@ -242,6 +296,8 @@ class GPUOptions(google.protobuf.message.Message):
result in undefined behavior.
"""
@property
def stream_merge_options(self) -> global___GPUOptions.Experimental.StreamMergeOptions: ...
def __init__(
self,
*,
@@ -260,8 +316,12 @@ class GPUOptions(google.protobuf.message.Message):
gpu_host_mem_limit_in_mb: builtins.float | None = ...,
gpu_host_mem_disallow_growth: builtins.bool | None = ...,
gpu_system_memory_size_in_mb: builtins.int | None = ...,
populate_pjrt_gpu_client_creation_info: builtins.bool | None = ...,
node_id: builtins.int | None = ...,
stream_merge_options: global___GPUOptions.Experimental.StreamMergeOptions | None = ...,
) -> None: ...
def ClearField(self, field_name: typing.Literal["collective_ring_order", b"collective_ring_order", "disallow_retry_on_allocation_failure", b"disallow_retry_on_allocation_failure", "gpu_host_mem_disallow_growth", b"gpu_host_mem_disallow_growth", "gpu_host_mem_limit_in_mb", b"gpu_host_mem_limit_in_mb", "gpu_system_memory_size_in_mb", b"gpu_system_memory_size_in_mb", "internal_fragmentation_fraction", b"internal_fragmentation_fraction", "kernel_tracker_max_bytes", b"kernel_tracker_max_bytes", "kernel_tracker_max_interval", b"kernel_tracker_max_interval", "kernel_tracker_max_pending", b"kernel_tracker_max_pending", "num_dev_to_dev_copy_streams", b"num_dev_to_dev_copy_streams", "num_virtual_devices_per_gpu", b"num_virtual_devices_per_gpu", "timestamped_allocator", b"timestamped_allocator", "use_cuda_malloc_async", b"use_cuda_malloc_async", "use_unified_memory", b"use_unified_memory", "virtual_devices", b"virtual_devices"]) -> None: ...
def HasField(self, field_name: typing.Literal["stream_merge_options", b"stream_merge_options"]) -> builtins.bool: ...
def ClearField(self, field_name: typing.Literal["collective_ring_order", b"collective_ring_order", "disallow_retry_on_allocation_failure", b"disallow_retry_on_allocation_failure", "gpu_host_mem_disallow_growth", b"gpu_host_mem_disallow_growth", "gpu_host_mem_limit_in_mb", b"gpu_host_mem_limit_in_mb", "gpu_system_memory_size_in_mb", b"gpu_system_memory_size_in_mb", "internal_fragmentation_fraction", b"internal_fragmentation_fraction", "kernel_tracker_max_bytes", b"kernel_tracker_max_bytes", "kernel_tracker_max_interval", b"kernel_tracker_max_interval", "kernel_tracker_max_pending", b"kernel_tracker_max_pending", "node_id", b"node_id", "num_dev_to_dev_copy_streams", b"num_dev_to_dev_copy_streams", "num_virtual_devices_per_gpu", b"num_virtual_devices_per_gpu", "populate_pjrt_gpu_client_creation_info", b"populate_pjrt_gpu_client_creation_info", "stream_merge_options", b"stream_merge_options", "timestamped_allocator", b"timestamped_allocator", "use_cuda_malloc_async", b"use_cuda_malloc_async", "use_unified_memory", b"use_unified_memory", "virtual_devices", b"virtual_devices"]) -> None: ...
PER_PROCESS_GPU_MEMORY_FRACTION_FIELD_NUMBER: builtins.int
ALLOW_GROWTH_FIELD_NUMBER: builtins.int
stubs/tensorflow/tensorflow/core/protobuf/rewriter_config_pb2.pyi
+2 -2
View File
@@ -333,8 +333,8 @@ class RewriterConfig(google.protobuf.message.Message):
(default is ON).
"""
auto_mixed_precision: global___RewriterConfig.Toggle.ValueType
"""Optimize data types for CUDA (default is OFF).
This will try to use float16 on GPU which is faster.
"""Optimize data types for CUDA/oneDNN (default is OFF).
This will try to use float16 on GPU/CPU which is faster.
Note that this can change the numerical stability of the graph and may
require the use of loss scaling to maintain model convergence.
"""
stubs/tensorflow/tensorflow/data/__init__.pyi
+2 -1
View File
@@ -114,6 +114,7 @@ class Dataset(ABC, Generic[_T1]):
element_spec: ContainerGeneric[tf.TypeSpec[Any]] | None = None,
compression: _CompressionTypes = None,
reader_func: Callable[[Dataset[Dataset[Any]]], Dataset[Any]] | None = None,
wait: bool = False,
) -> Dataset[Any]: ...
# PEP 646 could be used here for a more precise type when better supported.
def map(
@@ -196,7 +197,7 @@ class Dataset(ABC, Generic[_T1]):
self,
buffer_size: ScalarTensorCompatible,
seed: int | None = None,
reshuffle_each_iteration: bool | None = None,
reshuffle_each_iteration: bool = True,
name: str | None = None,
) -> Dataset[_T1]: ...
def skip(self, count: ScalarTensorCompatible, name: str | None = None) -> Dataset[_T1]: ...
stubs/tensorflow/tensorflow/saved_model/__init__.pyi
+5
View File
@@ -47,15 +47,19 @@ class SaveOptions:
"experimental_image_format",
"experimental_skip_saver",
"experimental_sharding_callback",
"extra_tags",
)
namespace_whitelist: list[str]
save_debug_info: bool
function_aliases: dict[str, PolymorphicFunction[..., object]]
experimental_debug_stripper: bool
experimental_io_device: str
experimental_variable_policy: VariablePolicy
experimental_custom_gradients: bool
experimental_image_format: bool
experimental_skip_saver: bool
experimental_sharding_callback: Incomplete | None
extra_tags: Incomplete | None
def __init__(
self,
namespace_whitelist: list[str] | None = None,
@@ -68,6 +72,7 @@ class SaveOptions:
experimental_image_format: bool = False,
experimental_skip_saver: bool = False,
experimental_sharding_callback: Incomplete | None = None,
extra_tags: Incomplete | None = None,
) -> None: ...
def contains_saved_model(export_dir: str | Path) -> bool: ...
stubs/tensorflow/tensorflow/tsl/protobuf/coordination_service_pb2.pyi
+4 -1
View File
@@ -473,15 +473,18 @@ class InsertKeyValueRequest(google.protobuf.message.Message):
DESCRIPTOR: google.protobuf.descriptor.Descriptor
KV_FIELD_NUMBER: builtins.int
ALLOW_OVERWRITE_FIELD_NUMBER: builtins.int
allow_overwrite: builtins.bool
@property
def kv(self) -> global___KeyValueEntry: ...
def __init__(
self,
*,
kv: global___KeyValueEntry | None = ...,
allow_overwrite: builtins.bool | None = ...,
) -> None: ...
def HasField(self, field_name: typing.Literal["kv", b"kv"]) -> builtins.bool: ...
def ClearField(self, field_name: typing.Literal["kv", b"kv"]) -> None: ...
def ClearField(self, field_name: typing.Literal["allow_overwrite", b"allow_overwrite", "kv", b"kv"]) -> None: ...
global___InsertKeyValueRequest = InsertKeyValueRequest
stubs/tensorflow/tensorflow/tsl/protobuf/dnn_pb2.pyi
+40
View File
@@ -267,6 +267,24 @@ FORWARD_BIAS_ACTIVATION: ConvolutionKind.ValueType # 4
FORWARD_GRAPH: ConvolutionKind.ValueType # 5
global___ConvolutionKind = ConvolutionKind
class _NormKind:
ValueType = typing.NewType("ValueType", builtins.int)
V: typing_extensions.TypeAlias = ValueType
class _NormKindEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_NormKind.ValueType], builtins.type):
DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
LAYER_FWD_INFER: _NormKind.ValueType # 0
LAYER_FWD_TRAIN: _NormKind.ValueType # 1
LAYER_BWD: _NormKind.ValueType # 2
class NormKind(_NormKind, metaclass=_NormKindEnumTypeWrapper):
"""NormKind kind"""
LAYER_FWD_INFER: NormKind.ValueType # 0
LAYER_FWD_TRAIN: NormKind.ValueType # 1
LAYER_BWD: NormKind.ValueType # 2
global___NormKind = NormKind
class _FusedMHAKind:
ValueType = typing.NewType("ValueType", builtins.int)
V: typing_extensions.TypeAlias = ValueType
@@ -285,6 +303,28 @@ BMM1_OUTPUT_INPUT_TYPE: FusedMHAKind.ValueType # 1
BMM1_OUTPUT_FLOAT: FusedMHAKind.ValueType # 2
global___FusedMHAKind = FusedMHAKind
class _FMHAMaskKind:
ValueType = typing.NewType("ValueType", builtins.int)
V: typing_extensions.TypeAlias = ValueType
class _FMHAMaskKindEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_FMHAMaskKind.ValueType], builtins.type):
DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
NO_MASK: _FMHAMaskKind.ValueType # 0
PADDING: _FMHAMaskKind.ValueType # 1
CAUSAL: _FMHAMaskKind.ValueType # 2
PADDING_CAUSAL: _FMHAMaskKind.ValueType # 3
ALIBI: _FMHAMaskKind.ValueType # 4
class FMHAMaskKind(_FMHAMaskKind, metaclass=_FMHAMaskKindEnumTypeWrapper):
"""FusedMHAMaskKind kind"""
NO_MASK: FMHAMaskKind.ValueType # 0
PADDING: FMHAMaskKind.ValueType # 1
CAUSAL: FMHAMaskKind.ValueType # 2
PADDING_CAUSAL: FMHAMaskKind.ValueType # 3
ALIBI: FMHAMaskKind.ValueType # 4
global___FMHAMaskKind = FMHAMaskKind
@typing.final
class TensorDescriptorProto(google.protobuf.message.Message):
"""Generic tensor representation."""