package xla
Get desktop application:
View/edit binary Protocol Buffers messages
rpc Connect (, )
protocol.proto:114Connects a node to the distributed coordinator node. Blocks until all tasks have connected. The service receives the number of nodes to expect as an option passed to its constructor.
message
protocol.proto:53int32 protocol_version = 1
int32 timeout_milliseconds = 2
int32 node_id = 3
We assume that each node knows its globally-unique node ID, provided by whatever mechanism launches the tasks. Node IDs should form a dense range of integers [0, num_nodes).
uint64 client_id = 4
A unique ID number for the client.
message
protocol.proto:66uint64 session_id = 1
rpc EnumerateDevices (, )
protocol.proto:119Blocking enumeration of devices, used by the GPU backend only. In parallel, all clients call EnumerateDevices() with their local device topology, and receive back a global topology in response.
message
protocol.proto:70uint64 session_id = 1
optional local_topology = 3
message
protocol.proto:75optional global_topology = 1
rpc Heartbeat (, )
protocol.proto:125Health-checking RPC. Workers send heartbeats to the coordinator at regular intervals. If the worker does not hear from the coordinator or the coordinator does not hear from the tasks, the tasks abort.
message
protocol.proto:98uint64 session_id = 1
int32 node_id = 2
message
protocol.proto:102(message has no fields)
rpc KeyValueGet (, )
protocol.proto:138Looks up a key in the key-value service. Blocks until the key is present or until `timeout` expires.
message
protocol.proto:79uint64 session_id = 1
bytes key = 2
int32 timeout_milliseconds = 3
message
protocol.proto:85bool found = 1
bytes value = 2
rpc KeyValueSet (, )
protocol.proto:141Updates the value associated with a key.
message
protocol.proto:90uint64 session_id = 1
bytes key = 2
bytes value = 3
message
protocol.proto:96(message has no fields)
rpc Shutdown (, )
protocol.proto:130Shutdown RPC. Workers send this RPC when they are ready to shut down; the RPC blocks until all tasks have indicated they are ready to shut down, or a timeout is reached.
message
protocol.proto:104uint64 session_id = 1
int32 node_id = 2
message
protocol.proto:108(message has no fields)
service XlaService
xla_service.proto:49/////////////////////// Global data requests
rpc Compile (, )
xla_service.proto:132Compiles the provided computation into executable. Returns the handle of the executable.
message
xla.proto:498optional computation = 1
The graph to be compiled.
optional execution_options = 2
Options that affect how XLA compiles code to service this request.
repeated input_shape_with_layout = 3
The layouts of the input arguments. If not set, the default layout will be used. Although the real arguments are not needed in compilation, the layouts of the arguments can affect the compilation.
message
xla.proto:511optional handle = 1
The handle to the executable.
Computes the value of a constant expression. The request contains the computation graph for the constant expression.
message
xla.proto:556optional computation = 1
optional output_layout = 2
message
xla.proto:561optional literal = 1
A LiteralProto is returned directly for this request.
Creates a channel handle that can be used to transfer data between two computations via a pair of Send and Recv instructions.
message
xla.proto:484ChannelHandle.ChannelType channel_type = 1
message
xla.proto:488optional channel = 1
rpc DeconstructTuple (, )
xla_service.proto:62Deconstructs a tuple. Returns a newly created GlobalDataHandle for each element in the tuple.
message
xla.proto:566optional tuple_handle = 2
message
xla.proto:570repeated element_handles = 1
rpc Execute (, )
xla_service.proto:137Invokes the provided executable with the provided global data passed as immutable arguments. The request contains the handle to the executable. Returns global data output and execution timing.
message
xla.proto:516optional handle = 1
repeated arguments = 2
The shape and layout of the arguments must be the same as the those of the executable's parameters.
Invokes the provided list of computations in parallel with the provided global data for each computation. Returns a list of global data output and execution timing.
message
xla.proto:534repeated requests = 1
message
xla.proto:543repeated responses = 1
Requests the statistics of the given computation.
message
xla.proto:475optional computation = 1
optional debug_options = 2
message
xla.proto:480optional stats = 1
Requests one or more device handles from the target. The returned device handles can be used to specify the device on which to execute computations or transfer data.
message
xla.proto:419int64 device_count = 1
message
xla.proto:423repeated device_handles = 1
rpc GetShape (, )
xla_service.proto:74Requests the shape of the referenced global data.
message
xla.proto:604optional data = 1
message
xla.proto:608optional shape = 1
rpc LoadData (, )
xla_service.proto:83Loads a variable number of values with a given element type from ColumnIO.
message
xla.proto:574string columnio_tablet_path = 1
Describes the path of the ColumnIO tablet to load.
string columnio_field = 2
Describes the field to load within the ColumnIO tablet.
optional element_shape = 3
Individual element shape, excluding rows.
int64 offset = 4
Warning: ColumnIO does not support random-access, so use offset with caution in performance-critical scenarios.
int64 limit = 5
Maximum number of elements (with shape element_shape) to load.
bool zip = 6
If more than one item is requested (via limit > 1), then this request attribute zips together the produced vectors.
message
xla.proto:596optional data = 1
optional data_shape = 2
int64 available_rows = 3
int64 rows_loaded = 4
int64 nanoseconds = 5
rpc ResetDevice (, )
xla_service.proto:108Resets the device, clearing all existing state on the device.
message
xla.proto:469optional device_handle = 1
message
xla.proto:473(message has no fields)
Transferred literal from the Outfeed buffer of the device.
message
xla.proto:456optional shape_with_layout = 1
This optional field directs the service to return the literal in this layout. A shape is used to hold the layout to accommodate tuples.
int64 replica_id = 2
optional device_handle = 3
message
xla.proto:465optional literal = 1
rpc TransferToClient (, )
xla_service.proto:87Transfers the given global data to the client in the form of a Literal.
message
xla.proto:427optional data = 1
optional shape_with_layout = 2
This optional field directs the service to return the literal in this layout. A shape is used to hold the layout to accommodate tuples.
message
xla.proto:435optional literal = 1
rpc TransferToInfeed (, )
xla_service.proto:98Transfers the given literal to the Infeed buffer of the device.
message
xla.proto:448optional literal = 1
int64 replica_id = 2
optional device_handle = 3
message
xla.proto:454(message has no fields)
rpc TransferToServer (, )
xla_service.proto:93Transfers the given literal to the server to be stored in a global allocation, which is returned.
message
xla.proto:439optional literal = 1
optional device_handle = 2
message
xla.proto:444optional data = 1
rpc Unpack (, )
xla_service.proto:70Unpack requests that a global data handle, with a tuple shape, has global data handles created for each of its constituent members. This is the equivalent of the "destructuring assignment" present in various programming languages.
message
xla.proto:612optional data = 1
message
xla.proto:616repeated tied_data = 1
rpc Unregister (, )
xla_service.proto:57Unregisters a global allocation. If the handle given is not currently allocated, a NOT_FOUND status is returned.
message
xla.proto:492repeated data = 1
message
xla.proto:496(message has no fields)
Waits until the given execution (aysnchronously launched) is complete, and returns the global data output.
message
xla.proto:547optional execution = 1
message
xla.proto:551optional output = 1
optional profile = 2
message
hlo.proto:440Serialization of BufferAllocation.
Used in:
int64 index = 1
int64 size = 2
bool is_thread_local = 3
bool is_tuple = 11
bool is_entry_computation_parameter = 5
bool is_constant = 12
int64 parameter_number = 6
repeated int64 parameter_shape_index = 10
bool maybe_live_out = 7
int64 color = 8
repeated assigned = 9
message
hlo.proto:443Assigned represents a single LogicalBuffer that is assigned to this BufferAllocation.
Used in:
int64 logical_buffer_id = 1
int64 offset = 2
int64 size = 3
message
hlo.proto:506Serialization of BufferAssignment.
Used in:
repeated logical_buffers = 1
repeated buffer_aliases = 2
repeated buffer_allocations = 3
repeated heap_simulator_traces = 4
message
hlo.proto:509Alias represents a source LogicalBuffer, and the buffer location that aliases it.
Used in:
int64 source_buffer_id = 1
optional location = 2
message
xla_data.proto:360Handle given to a user to represent a channel between two computations via a Send and Recv instruction pair. Channels are unbuffered, so Send Send instructions will be blocked until the data is transferred.
Used in:
int64 handle = 1
ChannelHandle.ChannelType type = 2
Used in: ,
CHANNEL_TYPE_INVALID = 0
Invalid primitive type to serve as default.
DEVICE_TO_DEVICE = 1
A channel for sending data between devices.
DEVICE_TO_HOST = 2
A channel for sending data from the device to the host. Can only be used with a Send operation.
HOST_TO_DEVICE = 3
A channel for sending data from the host to the device. Can only be used with a Recv operation.
message
xla_data.proto:623Used in:
bool lower = 1
If true, uses the lower triangle of `a`. If false, uses the upper triangle of `a`.
message
xla_data.proto:241Statistics of a computation.
Used in:
double flop_count = 1
The number of floating point operations in the computation.
double transcendental_count = 2
The number of transcendental operations (e.g., exp) in the computation.
message
xla_data.proto:520Used in: ,
int64 input_batch_dimension = 7
The number of the dimension that represents batch in the input.
int64 input_feature_dimension = 8
The number of the dimension that represents features in the input.
repeated int64 input_spatial_dimensions = 11
The dimension numbers for the spatial dimensions that the window moves through in the input.
int64 kernel_input_feature_dimension = 3
The number of the dimension that represents input features in the convolutional kernel (rhs).
int64 kernel_output_feature_dimension = 4
The number of the dimension that represents output features in the convolutional kernel (rhs).
repeated int64 kernel_spatial_dimensions = 6
The dimension numbers for the spatial dimensions that the window moves through in the kernel (rhs). window.strides(0) is the stride in the kernel_spatial_dimensions(0) dimension.
int64 output_batch_dimension = 9
The number of the dimension that represents batch in the output.
int64 output_feature_dimension = 10
The number of the dimension that represents features in the output.
repeated int64 output_spatial_dimensions = 12
The dimension numbers for the spatial dimensions that the window moves through in the output.
message
hlo.proto:383Used in:
int64 parameter = 1
repeated int64 index = 2
message
xla_data.proto:740Specifies a pair of output/operand buffers for kCustomCall that alias each other.
Used in:
repeated int64 output_shape_index = 1
int64 operand_index = 2
repeated int64 operand_shape_index = 3
Used in:
SCHEDULE_NONE = 0
SCHEDULE_LATEST = 1
SCHEDULE_EARLIEST = 2
message
xla.proto:25Debugging options for XLA. These options may change at any time - there are no guarantees about backward or forward compatibility for these fields.
Used in: , ,
bool xla_hlo_graph_addresses = 2
Show addresses of HLO ops in graph dump.
bool xla_hlo_profile = 9
Instrument the computation to collect per-HLO cycle counts.
repeated string xla_disable_hlo_passes = 30
List of HLO passes to disable/enable. These names must exactly match the pass names as specified by the HloPassInterface::name() method. At least one of xla_disable_hlo_passes and xla_enable_hlo_passes_only must be empty.
repeated string xla_enable_hlo_passes_only = 124
bool xla_disable_all_hlo_passes = 104
Disables all HLO passes. Notes that some passes are necessary for correctness and the invariants that must be satisfied by "fully optimized" HLO are different for different devices and may change over time. The only "guarantee", such as it is, is that if you compile XLA and dump the optimized HLO for some graph, you should be able to run it again on the same device with the same build of XLA.
int32 xla_backend_optimization_level = 31
Numerical optimization level for the XLA compiler backend; the specific interpretation of this value is left to the backends.
bool xla_embed_ir_in_executable = 33
Embed the compiler IR as a string in the executable.
bool xla_eliminate_hlo_implicit_broadcast = 35
Eliminate implicit broadcasts when lowering user computations to HLO instructions; use explicit broadcast instead.
bool xla_cpu_multi_thread_eigen = 60
When generating calls to Eigen in the CPU backend, use multi-threaded Eigen mode.
string xla_gpu_cuda_data_dir = 61
Path to directory with cuda/ptx tools and libraries.
bool xla_gpu_ftz = 62
Enable flush-to-zero semantics in the GPU backend.
bool xla_gpu_disable_multi_streaming = 63
Disable multi-streaming in the GPU backend.
bool xla_gpu_use_random_streams = 134
Debugging feature: if enabled, the GPU backend will assign HLO operators to randomly chosen streams. This is intended to trigger concurrency bugs.
bool xla_llvm_enable_alias_scope_metadata = 70
If true, in LLVM-based backends, emit !alias.scope metadata in generated IR.
bool xla_llvm_enable_noalias_metadata = 71
If true, in LLVM-based backends, emit !noalias metadata in the generated IR.
bool xla_llvm_enable_invariant_load_metadata = 72
If true, in LLVM-based backends, emit !invariant.load metadata in the generated IR.
bool xla_llvm_disable_expensive_passes = 73
If true, a set of expensive LLVM optimization passes will not be run.
bool xla_test_all_output_layouts = 90
This is used by ClientLibraryTestBase::ComputeAndCompare*. If true, the computation will run n! times with all permunations of layouts for the output shape in rank n. For example, with a 3D shape, all permutations of the set {0, 1, 2} are tried.
bool xla_test_all_input_layouts = 91
This is used by ClientLibraryTestBase::ComputeAndCompare*. If true, the computation will run for all permunations of layouts of all input arguments. For example, with 2 input arguments in 2D and 4D shapes, the computation will run 2! * 4! times.
bool xla_hlo_graph_sharding_color = 92
Assign colors based on sharding information when generating the Graphviz HLO graph.
bool xla_gpu_use_cudnn_batchnorm = 94
If true, the GPU backend is free to use cudnn for HLO batch normalization ops.
bool xla_cpu_use_mkl_dnn = 97
Generate calls to MKL-DNN in the CPU backend.
int32 xla_gpu_max_kernel_unroll_factor = 98
Maximum kernel unroll factor for the GPU backend.
bool xla_cpu_enable_fast_math = 99
When true, "unsafe" mathematical optimizations are enabled. These transformations include but are not limited to: - Reducing the precision of operations (e.g. using an approximate sin function, or transforming x/y into x * (1/y)). - Assuming that operations never produce or consume NaN or +/- Inf (this behavior can be adjusted using xla_cpu_fast_math_allow_{nans|infs}). - Assuming that +0 and -0 are indistinguishable.
bool xla_cpu_fast_math_honor_nans = 120
When xla_cpu_enable_fast_math is true then this controls whether we allow operations to produce NaNs. Ignored when xla_cpu_enable_fast_math is false.
bool xla_cpu_fast_math_honor_infs = 121
When xla_cpu_enable_fast_math is true then this controls whether we allow operations to produce infinites. Ignored when xla_cpu_enable_fast_math is false.
bool xla_cpu_fast_math_honor_division = 126
When xla_cpu_enable_fast_math is true then this controls whether we forbid to use the reciprocal of an argument instead of division. Ignored when xla_cpu_enable_fast_math is false.
bool xla_cpu_fast_math_honor_functions = 129
When xla_cpu_enable_fast_math is true then this controls whether we forbid to approximate calculations for functions. Ignored when xla_cpu_enable_fast_math is false.
bool xla_cpu_enable_fast_min_max = 140
When false we lower the Minimum and Maximum hlos in the CPU backend such that Min(NotNaN, NaN) = Min(NaN, NotNaN) = NaN. In other words, if flag this is false we always propagate NaNs through Min and Max. Note, this does not correspond to the exact same behavior as the gpu flag below!
bool xla_gpu_enable_fast_min_max = 100
When true we lower the Minimum and Maximum hlos in the GPU backend such that Min(NotNaN, NaN) = Min(NaN, NotNaN) = NotNaN. In other words, if flag this is true we don't propagate NaNs through Min and Max. Note, this does not correspond to the exact same behavior as the cpu flag above!
bool xla_allow_excess_precision = 122
Allows xla to increase the output precision of floating point operations.
bool xla_gpu_crash_on_verification_failures = 101
Crashes the program when any kind of verification fails, instead of just logging the failures. One example is cross checking of convolution results among different algorithms.
int32 xla_gpu_autotune_level = 123
0: Disable gemm and convolution autotuning. 1: Enable autotuning, but disable correctness checking. 2: Also set output buffers to random numbers during autotuning. 3: Also reset output buffers to random numbers after autotuning each algorithm. 4+: Also check for correct outputs and for out-of-bounds reads/writes. Default: 4.
int32 xla_force_host_platform_device_count = 102
Force the host platform to pretend that there are these many host "devices". All these devices are backed by the same threadpool. Defaults to 1. Setting this to anything other than 1 can increase overhead from context switching but we let the user override this behavior to help run tests on the host that run models in parallel across multiple devices.
bool xla_gpu_disable_gpuasm_optimizations = 103
If set to true XLA:GPU invokes `ptxas` with -O0 (default is -O3).
bool xla_hlo_evaluator_use_fast_path = 106
Enable fast math with eigen in the HLO evaluator.
bool xla_allow_scalar_index_dynamic_ops = 107
Temporary option to allow support for both the R1 and the scalar index versions of DynamicSlice and DynamicUpdateSlice. Only used for testing.
DebugOptions.StepMarkerLocation xla_step_marker_location = 108
Option to emit a target-specific marker to indicate the start of a training step. The location of the marker (if any) is determined by the option value.
string xla_dump_to = 109
Directory to dump into.
string xla_dump_hlo_module_re = 110
If specified, will only dump modules which match this regexp.
string xla_dump_hlo_pass_re = 111
If this flag is specified, will also HLO before and after passes that match this regular expression. Set to .* to dump before/after all passes.
bool xla_dump_hlo_as_text = 112
Specifies the format that HLO is dumped in. Multiple of these may be specified.
bool xla_dump_hlo_as_proto = 113
bool xla_dump_hlo_as_dot = 114
bool xla_dump_hlo_as_url = 115
bool xla_dump_hlo_as_html = 116
Dump HLO graphs as an HTML (DOT -> SVG inlined in HTML)
bool xla_dump_fusion_visualization = 149
Dump the visualization of the fusion progress.
bool xla_dump_hlo_snapshots = 118
If true, every time an HLO module is run, we will dump an HloSnapshot (essentially, a serialized module plus its inputs) to the --xla_dump_to directory.
bool xla_dump_include_timestamp = 131
Include a timestamp in the dumped filenames.
int32 xla_dump_max_hlo_modules = 132
Max number of hlo module dumps in a directory. Set to < 0 for unbounded.
bool xla_dump_module_metadata = 144
Dump HloModuleMetadata as a text proto for each HLO module.
bool xla_dump_compress_protos = 151
GZip-compress protos dumped via --xla_dump_hlo_as_proto.
bool xla_gpu_force_conv_nchw = 125
Overrides for XLA GPU's convolution layout heuristic.
bool xla_gpu_force_conv_nhwc = 146
repeated string xla_gpu_ptx_file = 127
Paths to files with ptx code.
string xla_gpu_algorithm_denylist_path = 128
Denylist for cuDNN convolutions.
bool xla_gpu_deterministic_reductions = 130
Guarantee run-to-run determinism from reductions on XLA:GPU.
bool xla_tpu_detect_nan = 135
Debug options that trigger execution errors when NaN or Inf are detected.
bool xla_tpu_detect_inf = 136
bool xla_cpu_enable_xprof_traceme = 137
True if TraceMe annotations are enabled for XLA:CPU.
bool xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found = 138
It is usually preferable to not fallback to the driver; it can consume more memory, or have bugs.
string xla_gpu_asm_extra_flags = 141
Extra parameters to pass the GPU assembler.
int32 xla_multiheap_size_constraint_per_heap = 142
Per-heap size constraint. New heaps will be created if per-heap max size is reached.
bool xla_detailed_logging_and_dumping = 143
Enable detailed logging into vlog and xla dumping. If this is disabled, no compilation summary will be printed in the end of computation and no hlo modules will be dumped.
int32 xla_gpu_force_compilation_parallelism = 147
Overrides normal multi-threaded compilation settting to use this many threads. Setting to 0 (the default value) means no enforcement.
bool xla_gpu_deterministic_ops = 148
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.
repeated string xla_gpu_llvm_ir_file = 150
Paths to files with LLVM code.
bool xla_gpu_enable_async_all_reduce = 152
Convert synchronous all-reduces ops into asynchronous.
map<string, string> xla_backend_extra_options = 500
Extra options to pass to the compilation backend (e.g. LLVM); specific interpretation of these values is left to the backend.
Used in: ,
STEP_MARK_AT_ENTRY = 0
Generate a step marker at the program entry. This handles the case where each step is done by one or multiple program execution(s). Only the first program will be tagged for generating a step marker at the program entry. This is the default.
STEP_MARK_AT_TOP_LEVEL_WHILE_LOOP = 1
Generate a step marker at each iteration of the top level while loop, which is assumed to be a training loop.
STEP_MARK_AT_SECOND_LEVEL_WHILE_LOOP = 3
Generate a step marker at each iteration of the second level while loops, which is assumed to be a training or eval loop.
STEP_MARK_NONE = 2
No step marker generated.
message
xla_data.proto:383DeviceAssignmentProto is a serialized form of DeviceAssignment class, which represents the device ids assigned to a set of replicated computations. See xla::DeviceAssignment class comment for more details.
Used in: , , ,
int32 replica_count = 1
int32 computation_count = 2
repeated computation_devices = 3
message
xla_data.proto:389Each logical computation runs on replica_count physical devices. ComputationDevice represents the device ids assinged to the replicas.
Used in:
repeated int32 replica_device_ids = 1
message
xla_data.proto:349Handle given to a user that represents a replicated virtual device. Each replicated device represents N physical devices for execution where N is the number of replicas.
Used in: , , , , ,
int64 handle = 1
int64 device_count = 2
The number of model-parallel virtual devices that communicate via XLA Send/Recv instructions.
message
protocol.proto:32Describes a device local to a host.
Used in:
int32 local_device_ordinal = 1
string name = 2
string vendor = 3
int32 global_device_id = 4
The following fields are present in the GlobalTopologyProto message returned by Connect() but not in the LocalTopologyProto messages passed to Connect(). In other words, the coordinator node determines the global device IDs during Connect().
Globally unique ID number.
message
xla_data.proto:571Used in: ,
repeated int64 lhs_contracting_dimensions = 1
The dimension numbers that represent the 'lhs' contracting dimensions.
repeated int64 rhs_contracting_dimensions = 2
The dimension numbers that represent the 'rhs' contracting dimensions.
repeated int64 lhs_batch_dimensions = 3
The dimension numbers that represent the 'lhs' batch dimensions.
repeated int64 rhs_batch_dimensions = 4
The dimension numbers that represent the 'rhs' batch dimensions.
message
hlo.proto:348Used in:
repeated entries = 1
message
hlo.proto:372A list of bindings which indicates that the `target_dim_num` in the subshape `target_param_index` of parameter `target_param_num` is a dynamic dimension and its real dynamic size is represented by `dynamic_param_index` in parameter `dynamic_param_num`. As an example, imagine we have a program: ENTRY main { a = f32[] parameter(0) b = f32[10] parameter(1) ROOT root = (f32[], f32[10]) tuple(%a, %b) } Let's say 'b' (param index 1) is a dynamic shape whose input has an upperbound of 10 and real size is determined at runtime.'a' represents the real size of b's first dimension. In this case, the fields are set in the following way: dynamic_param_num = 1 dynamic_param_index = {} target_param_num = 0 target_param_index = {} target_param_dim = 0
Used in:
int64 dynamic_param_num = 1
repeated int64 dynamic_param_index = 2
int64 target_param_num = 3
repeated int64 target_param_index = 4
int64 target_param_dim_num = 5
message
xla.proto:526TODO(b/118493728): Remove this and ExecuteGraphParallelRequest and replace the uses with calls to Compile and Execute.
Used in:
optional computation = 1
repeated arguments = 2
optional execution_options = 3
Options that affect how XLA compiles and runs code to service this request.
message
xla.proto:538Used as response type in: XlaService.Execute
Used as field type in:
optional output = 1
optional profile = 2
message
xla_data.proto:335Handle given to a user that represents an execution that the user launched asynchronously on the device.
Used in: , ,
int64 handle = 1
message
xla.proto:351These settings control how XLA compiles and/or runs code. Not all settings will have an effect on every platform. When adding new fields, keep in mind that boolean fields default to false.
Used in: ,
optional shape_with_output_layout = 2
This optional field's layout is used as a hint when storing the output of this computation. Subsequent transfers of this output array to the client may be faster when using this layout. We use a Shape here to accommodate computations that return a tuple.
uint64 seed = 3
Used to seed random-number generators used in this computation. If this is 0, we generate a seed ourselves. TODO(b/32083678): Changing the seed unnecessarily forces a recompilation.
optional debug_options = 4
repeated device_handles = 5
This optional field specifies a particular set of devices to run the computation on. The computation will be partitioned across these devices. If not provided, the default device will be chosen.
int32 num_replicas = 6
Number of replicas of the computation to run. If zero, uses the default number of replicas for the XLA service.
optional device_assignment = 7
This optional field specifies the device assignment if known at compile time.
bool alias_passthrough_params = 8
Alias input and output buffers for parameters that are passed-through XLA modules without being changed.
int32 num_partitions = 9
Number of partitions of the computation to run (model parallelism). If zero, uses the default number of partitions for the XLA service.
int32 launch_id = 10
Used to identify a set of programs that should be launch together.
bool use_spmd_partitioning = 11
Indicates whether to use SPMD (true) or MPMD (false) partitioning when num_partitions > 1 and XLA is requested to partition the input program.
bool deduplicate_hlo = 12
If set, deduplicate hlo into function calls to reduce binary size. Only works on TPU.
repeated int32 argument_input_indices = 100
IPU Specific. The indices of arguments to the original TF operation. This does not include compile time constants or resource variables.
repeated int32 resource_input_indices = 101
The indices of resources to the TF operation which are resource variables.
repeated bool resource_input_initialized = 102
Element at index `i` indicates whether resource variable at `resource_input_indices[i]` is initialized.
repeated int32 resource_update_to_input_index = 103
The number of outputs which are updates for resource variables. For each output which is a resource update, this contains a mapping to the input parameter which is for that resource.
message
xla_data.proto:302Profile data from the execution of a computation.
Used in: ,
bool compilation_cache_hit = 1
Whether the executable was read from the compilation cache.
int64 compile_time_ms = 2
The time in milliseconds spent to compile the computation. This only set if the executable was not read from the compilation cache (compilation_cache_hit == false).
int64 compute_cycle_count = 3
The number of cycles spent for the computation. This does not include the time taken for the data transfers between the host and the device. This is a target-dependent field and only used for debugging purposes.
int64 compute_time_ns = 4
The time in nanoseconds spent for the computation, without data transfer.
int64 compute_and_transfer_time_ns = 5
The time in nanoseconds spent for the entire computation, including the result data transfer time. Current implementation does not spend any cycles for the input data transfer since the memory is initialized with the proper values before the execution.
int64 executable_size_in_bytes = 6
The size of the binary code in the executable.
bool profile_cache_hit = 7
Whether this profile was drawn from a cache of profiles instead of from execution on the hardware.
Used in:
FFT = 0
Forward FFT; complex in, complex out.
IFFT = 1
Inverse FFT; complex in, complex out.
RFFT = 2
Forward real FFT; real in, fft_length / 2 + 1 complex out
IRFFT = 3
Inverse real FFT; fft_length / 2 + 1 complex in,
A format specifies the method used by a layout to store an array in memory.
Used in:
INVALID_FORMAT = 0
TODO(b/120869032): Rename this to FORMAT_NONE or something else which better corresponds to its meaning.
DENSE = 1
The default layout, with exactly one storage location per element.
message
xla_data.proto:631Generic map of attributes used to pass hints / configuration options from the Python frontend to the XLA backend.
Used in:
map<string, string> map = 1
message
xla_data.proto:479Describes the dimension numbers for a gather operation. See https://www.tensorflow.org/performance/xla/operation_semantics#gather for more details.
Used in:
repeated int64 offset_dims = 1
"Window indices" is a term for a set of indices that index into the interior of a dynamic-slice from the input tensor, the starting indices for which were computed from output_gather_dims (see the operation semantic for how this is defined) and the start_indices tensor. The window indices for a specific output index Out is computed as: i = 0 for (k : [0, input_tensor_shape.rank)) window_indices[k] = if k in collapsed_slice_dims then 0 else Out[offset_dims[i++]]
repeated int64 collapsed_slice_dims = 2
repeated int64 start_index_map = 3
This is interpreted as a map from i to start_index_map[i]. It transforms the gather index looked up from the start_indices tensor into the starting index in the input space.
int64 index_vector_dim = 4
The dimension in the start_indices input that contains the starting indices.
message
xla_data.proto:342Handle given to a user that represents a globally accessible allocation. Contrast this against a ComputationDataHandle, which is not globally accessible, since it only exists within a specific computation.
Used in: , , , , , , , , , , , ,
int64 handle = 1
message
protocol.proto:49Used in:
repeated nodes = 1
message
hlo.proto:463A trace of a HeapSimulator run.
Used in:
repeated events = 1
bool whole_module_simulation = 2
int64 buffer_allocation_index = 3
message
hlo.proto:466The trace includes a list of events, where each event describes one action performed by the heap simulator.
Used in:
Event.Kind kind = 1
int64 buffer_id = 2
The id of the LogicalBuffer that the event applies to.
string computation_name = 3
The HloInstruction that the simulation was processing that caused this event to occur, identified by its computation and instruction name. E.g. buffers defined by instruction A are allocated when processing A.
string instruction_name = 4
int64 share_with_canonical_id = 5
The id of the canonical LogicalBuffer that the buffer shares with. Only set for SHARE_WITH events.
Used in:
ALLOC = 0
A memory region was allocated for the buffer.
FREE = 1
A memory region was freed for the buffer.
SHARE_WITH = 2
A buffer was shared with another (canonical) buffer. This is similar to ALLOC, except that instead of allocating a new region of memory, the memory region of the canonical buffer is directly re-used. Multiple buffers may share with the same canonical buffer. The lifetime of the canonical buffer is extended to the union of all lifetimes.
message
hlo.proto:278Serialization of HloComputation.
Used in:
string name = 1
repeated instructions = 2
The array of instructions is always in a valid dependency order, where operands appear before their users.
optional program_shape = 4
int64 id = 5
The id of this computation.
int64 root_id = 6
The id of the root of the computation.
optional printer_data = 1
repeated int64 profile_counters = 2
message
hlo.proto:320Used in:
repeated entries = 1
message
hlo.proto:334The following proto describes a pair of aliased an input (described by parameter number and a ShapeIndex of the parameter) and an output (described by a ShapeIndex of the root instruction). For example: entry = { output_shape_index={1}, parameter_number=0, parameter_shape_index={1, 2}, } This entry indicates that the first paremter's {1, 2} element is aliased with the {1} element of the root instruction.
Used in:
repeated int64 output_shape_index = 1
ShapeIndex of the root hlo.
int64 parameter_number = 2
Number of the parameter in entry computation.
repeated int64 parameter_shape_index = 3
ShapeIndex of the parameter instruction.
Kind kind = 4
The kind of alias to be setup.
message
hlo.proto:45Serialization of HloInstruction. Next ID: 77
Used in: ,
string name = 1
string opcode = 2
optional shape = 3
optional metadata = 7
optional literal = 8
Literal, only present for kConstant.
int64 parameter_number = 9
Parameter number is only present for kParameter.
string fusion_kind = 11
Fusion state, only present for kFusion.
int64 tuple_index = 13
Index for kGetTupleElement.
repeated int64 dimensions = 14
Dimensions present for some operations that require reshaping or broadcasting, including Reshape, Reduce, ReduceWindow, and Reverse.
optional window = 15
Describes the window in a windowed operation such as convolution.
optional convolution_dimension_numbers = 16
Describes the dimension numbers used for a convolution.
int64 feature_group_count = 50
The number of feature groups. Used for a convolution. Must be a divisor of the input feature dimension and output feature dimension. If not specified, it will use a default value of 1.
int64 batch_group_count = 58
repeated slice_dimensions = 17
int32 exponent_bits = 18
The bit sizes for a reduce-precision operation.
int32 mantissa_bits = 19
repeated int64 dynamic_slice_sizes = 20
Describes the [start, start + size) range size for a dynamic slice ('start' is specified dynamically in the second operand of the operation).
optional padding_config = 21
The padding configuration that describes the edge padding and interior padding of this pad instruction. Only set for pad instructions.
bytes outfeed_config = 22
Outfeed configuration information, only present for kOutfeed.
RandomDistribution distribution = 23
The distribution requested for random number generation. Only present for kRng.
float epsilon = 24
A small float number added to the variance to avoid divide-by-zero error. Only present for kBatchNormTraining.
int64 feature_index = 25
An integer value representing the index of the feature dimension. Only present for kBatchNormTraining.
int64 channel_id = 26
Represents a unique identifier for each Send/Recv instruction pair or optionally for collective instructions (AllReduce, CollectivePermute, AllToAll). Non-positive channel_id is equivalent to no channel id.
bytes infeed_config = 27
The string representation of the infeed configuration.
string custom_call_target = 28
Name of a external target (eg, global symbol) to call, only present for kCustomCall.
optional outfeed_shape = 29
Shape of outfeed request.
optional dot_dimension_numbers = 30
Describes the dimension numbers used for a dot operation
FftType fft_type = 31
FFT type (FFT, IFFT, etc).
repeated int64 fft_length = 32
FFT length.
string comparison_direction = 63
Comparison direction only used for kCompare.
optional gather_dimension_numbers = 33
Gather dimension numbers.
repeated int64 gather_slice_sizes = 34
string channel_name = 41
Compute Host.
int64 cost_estimate_ns = 42
int64 id = 35
The id of this instruction.
repeated int64 operand_ids = 36
repeated int64 control_predecessor_ids = 37
repeated int64 called_computation_ids = 38
optional sharding = 40
bytes backend_config = 43
Backend configuration for the instruction. Has backend-specific meaning.
repeated replica_groups = 49
Cross replica op fields.
int64 all_reduce_id = 45
Deprecated, but keeping it for backward compatibility. Use channel_id. Non-positive all_reduce_id is equivalent to no all_reduce_id.
bool use_global_device_ids = 71
If true, interprets ids in ReplicaGroup as global device ids, which is a linearized id of `replica_id * partition_count + partition_id`.
bool is_host_transfer = 47
Whether this Send/Recv instruction transfers data to/from the host. Only present for Send and Recv instructions and their SendDone and RecvDone partners.
bool is_stable = 60
Whether this Sort instruction should be stable.
optional scatter_dimension_numbers = 48
optional precision_config = 51
Precision configuration for the instruction. Has backend-specific meaning.
repeated source_target_pairs = 52
Collective permute field.
optional domain_entry_sharding = 54
Sharding for kDomain instructions.
optional domain_exit_sharding = 55
bool constrain_layout = 56
For custom call this indicates that the layouts are constrained. If constrain_layout is true then the 'shape' field must contain a layout, and 'operand_shapes_with_layout' must contain a shape with layout for each operand.
repeated operand_shapes_with_layout = 57
optional triangular_solve_options = 59
Options for TriangularSolve
optional cholesky_options = 62
Options for Cholesky
optional parameter_replication = 61
Describes how parameters behave with regards to replicas.
repeated int64 outer_dimension_partitions = 64
If set, the given instruction is run in parallel on e.g. multiple CPU cores. The outermost dimension gets split up into outer_dimension_partitions[0] pieces, the next-outermost dim gets split into outer_dimension_partitions[1] pieces, etc. It's illegal to partition a dimension into more shards than there are elements in that dimension.
bool custom_call_has_side_effect = 65
Whether the kCustomCall instruction has side-effects, only present for kCustomCall.
repeated custom_call_output_operand_aliasing = 74
A list of CustomCallOutputOperandAliasing pairs that specifies aliasing buffers between output and operands for kCustomCall.
CustomCallSchedule custom_call_schedule = 76
Specifies the desired schedule for the custom-call. The field is only present for custom-call.
int64 delta = 66
The delta value for kRngGetAndUpdateState.
bool indices_are_sorted = 67
Specifies if the gather/scatter indices are guaranteed to be sorted by the caller.
optional frontend_attributes = 68
Frontend attributes to pass to the XLA backend.
bool unique_indices = 69
Specifies if all elements updated are guaranteed to be unique by the caller.
RandomAlgorithm rng_algorithm = 70
RNG algorithm used by kRngBitGenerator.
string comparison_type = 72
The comparison type used for kCompare.
bool is_cross_program_prefetch = 73
Specifies if this is a cross-program-prefetch, used by kCopyStart.
PaddingType padding_type = 75
If a convolution is dynamic, a dynamic padding type will be specified.
message
hlo.proto:101Describes the [begin, end) index range and stride for slices.
Used in:
int64 start = 1
int64 limit = 2
int64 stride = 3
message
hlo.proto:500An abstraction representing a set of HLO module built to run concurrently across different devices.
string name = 1
repeated hlo_modules = 2
message
hlo.proto:549Metadata for an HLO module. Dumped after HLO passes and before LLO lowering with filename module_####.metadata.textproto, where #### is canonical_module_id.
int64 canonical_module_id = 1
Uniquely identifies an HloModuleMetadata. Equal to the first unique_id of the module (a module may go through multiple unique_ids). If a module is partitioned into multiple modules, those modules will each have a new HloModuleMetadata with a different canonical_module_id.
string module_group_name = 2
Name of the module group that the module is part of.
int64 original_module_id = 3
The canonical module id of the module that this one is partitioned from, if applicable.
repeated int64 partitioned_module_ids = 4
The canonical module ids of the modules that this one is partitioned into, if applicable.
repeated pass_metadata = 5
Metadata for the HLO passes that are run on the module.
message
hlo.proto:389Serialization of HloModule.
Used in: , , , , , ,
string name = 1
string entry_computation_name = 2
int64 entry_computation_id = 6
repeated computations = 3
The array of computations is always in a valid dependency order, where callees appear before their callers.
optional host_program_shape = 4
The host program shape (with layout) of the entry computation.
int64 id = 5
The id of this module.
optional schedule = 7
The schedule for this module.
optional input_output_alias = 8
Describes alias information between inputs and outputs.
optional dynamic_parameter_binding = 9
repeated cross_program_prefetches = 10
bool is_dynamic = 11
True if the module contains dynamic computation.
message
hlo.proto:576Metadata for one run of an HLO pass on a module. Provides more information when processing debug dumps of HloProtos about the order of HLO passes and various other stats like duration. `pass_id` may also be used to identify a particular run of a pass in debug info that propagates through stages of compilation.
Used in:
int64 pass_id = 1
For a given module, pass_id uniquely identifies a run of an HLO pass on that module. Note that a pass_id may not always refer to the same pass because the order of passes during compilation may change. For finding metadata for a particular pass, pass_name and pipeline_name would be more reliable, although note that they may not be unique.
string pass_name = 2
string pipeline_name = 3
repeated string dump_filenames = 4
Filenames of the dumps of the module after this pass ran. Module may be dumped in multiple formats, and the order of formats in this field will stay consistent across passes.
bool module_changed = 5
Return value of pass.Run(). True if this pass changed the module, or, in the case where the module was run through this pass as part of a module group, true if this pass changed any module in the same module group.
int64 module_id = 6
The unique_id of the module that this pass is run on. May be different from the canonical_module_id of the HloModuleMetadata that this HloPassMetadata is inside.
repeated int64 module_group_module_ids = 7
If the module went through this pass as part of a module group, this is set as the ids of all the modules in the module group. Empty otherwise.
int64 start_timestamp_usec = 8
Timestamp before and after the pass is run. Note they may be equal.
int64 end_timestamp_usec = 9
Describes how to pretty-print a profile counter array gathered for a specific HloModule.
Used in:
repeated computation_infos = 1
HloComputationInfos for every HloComputation in the HloModule.
int64 profile_counters_size = 2
The size of the profile counters array we will pretty-print.
map<string, int64> extra_metrics = 3
Maps extra metric name to the index into the profile counters array.
string entry_computation = 4
Name of the entry computation.
Pretty-printer information about an HloComputation.
Used in:
string name = 1
int64 profile_index = 2
The index into the profile counters array for the HloComputation corresponding to this HloComputationInfo.
repeated instruction_infos = 3
HloInstructionInfos for every HloInstruction in the HloComputation for corresponding to this HloComputattionInfo.
Pretty-printer information about an HloInstruction.
Used in:
string long_name = 1
string short_name = 2
string category = 3
float flop_count = 4
Metrics computed by HloCostAnalysis.
float transcendental_count = 5
int64 bytes_accessed = 9
float optimal_seconds = 7
int64 profile_index = 8
The index into the profile counters array for the HloInstruction corresponding to this HloInstructionInfo.
message
hlo.proto:521Grouping message that contains all of the information above.
Used in: , ,
optional hlo_module = 1
optional buffer_assignment = 3
message
hlo.proto:301Serialization of an HLO schedule. An HLO schedule contains a total order of instructions for each non-fusion computation in the module.
Used in:
map<int64, > sequences = 1
Map from computation id to sequence.
message
hlo.proto:302Used in:
repeated int64 instruction_ids = 1
message
hlo.proto:532Encapsulates HloProto together with the arguments, result, and execution_platform. This message is used for purposes such as analysis/replay/file-storage.
Used in: ,
optional hlo = 1
The hlo graph.
repeated arguments = 2
The arguments passed to the graph.
optional result = 3
The result of the graph.
string execution_platform = 4
The name of the platform used to run the graph.
enum Kind
hlo.proto:310Used in:
UNDEFINED_ALIAS = 0
Define a UNDEFINED_ALIAS equal to zero to get around the default-0 proto3 behavior and missing has_*() APIs.
MAY_ALIAS = 1
The buffers may or may not alias at runtime.
MUST_ALIAS = 2
The buffers must alias at runtime.
message
xla_data.proto:144A layout describes how the array is placed in (1D) memory space. This includes the minor-to-major ordering of dimensions within a shape. Clients must specify the layouts of input Literals to the computation. Layouts specified in interior operations which take Shapes (for example, Convert) are ignored. See the XLA documentation for more information on shapes and layouts. LINT.IfChange
Used in: , ,
Format format = 4
The method used to store the data in memory. The format determines which of the other fields are used by the layout.
repeated int64 minor_to_major = 1
Sequence of dimension numbers, from minor (fastest varying index) to major (slowest varying index). This field is required.
repeated tiles = 6
A sequence of tiles, starting from the tile that's applied first to the Shape. TODO(b/119839262): implement tiling in each backend or add Unimplemented error.
int64 element_size_in_bits = 7
Bit size of each element. If the size is bigger than what the element type requires, the value is stored in the least significant bits and the additional most significant bits are filled with 0's. TODO(b/119839262): implement in each backend or add Unimplemented error.
int64 memory_space = 8
Memory space where this array resides. The integer field is interpreted in a backend-specific manner.
message
xla_data.proto:401Literals are used when the server and client need to exchange materialized data / results. Literals are also used to describe constants used in computations. Transfers to/from the client are encoded in literal form, and the structure of the repeated fields is implied by the shape.
Used in: , , , , , , ,
optional shape = 1
repeated bool preds = 2
bytes s8s = 15
bytes u8s = 3
repeated int32 s32s = 4
repeated int64 s64s = 5
repeated uint32 u32s = 6
repeated uint64 u64s = 7
repeated float f32s = 8
repeated double f64s = 9
repeated float c64s = 12
Stored as interleaved real, imag floats.
repeated double c128s = 18
Stored as interleaved real, imag doubles.
repeated tuple_literals = 10
bytes f16s = 11
The F16s, BF16s, U16s and S16s are encoded in little endian byte order
bytes bf16s = 13
bytes u16s = 16
bytes s16s = 17
repeated int64 sparse_indices = 14
Next = 19
message
protocol.proto:44Used in: ,
int32 node_id = 1
repeated devices = 2
message
hlo.proto:419Serialization of LogicalBuffer.
Used in:
int64 id = 1
int64 size = 2
optional defined_at = 3
The location where the buffer is defined.
int64 color = 4
message
hlo.proto:422Location represents an instruction and its shape index, which uniquely identifies a point where a buffer is needed.
Used in: ,
string computation_name = 1
NOTE: module_name isn't necessary, since all LogicalBuffers are associated with a single HloModule.
string instruction_name = 2
repeated int64 shape_index = 3
message
xla_data.proto:261Symbolization metadata for HLO Instructions. This metadata is used for debugging XLA code generation, as well as performance profiling of XLA-generated executables.
Used in: ,
string op_type = 1
The framework op name that generated this XLA op. Frameworks that build on top of XLA should mirror the names of their ops back to users by specifying the op_type. In this way, even if the framework's "ops" are implemented as multiple XLA HLO Ops, they can be grouped appropriately. (e.g. if a SoftMax layer is emitted into XLA as multiple ops, then each op should have the op_type be "SoftMax".)
string op_name = 2
The user-specified name of the op. This name is often unique within a computation. Note: some frameworks add auto-generated names if the user does not provide one.
string source_file = 3
Indicate a file and line that this op is associated to in a user's program. e.g. it could be the file and line of user code that generated the op.
int32 source_line = 4
repeated ProfileType profile_type = 5
int64 creation_pass_id = 6
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.
int64 logical_creation_pass_id = 7
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.
int64 size_of_generated_code_in_bytes = 8
The footprint of the generated code for the instruction.
int64 size_of_memory_working_set_in_bytes = 9
The size of the working set, i.e., the amount of memory, used by the instruction in a compiler-managed fast device memory.
message
xla_data.proto:635Used in: , ,
OpSharding.Type type = 1
optional tile_shape = 2
The shape of the sharded tile.
repeated int64 tile_assignment_dimensions = 3
The shape of the tile assignment tensor - this must be the same rank as tile_shape and the product of its dimensions must equal tile_assignment_devices.size().
repeated int64 tile_assignment_devices = 4
Flattened list of device IDs. The order of flattening is the same as used by IndexUtil::MultiToLinearIndex(tile_assignment_shape).
repeated tuple_shardings = 5
If type == TUPLE, the sub-shardings, one per leaf node in the tuple shape, in pre-order. The tuple shape could be nested; here we store just a flattened list of all leaves in the tuple shape. Note that the tuple shape is not stored here; shardings do not store the shapes to which they are applied, this is inferred from the instruction this sharding gets attached to.
bool replicate_on_last_tile_dim = 6
Only used for OTHER type. If true, data is sharded according to other dimensions of tile_assignment(), but replicated across devices along the last dimension. (Experimental)
repeated metadata = 7
This field is used to track the source of this sharding, usually derived from instructions. Multple metadata may be populated if sharding is combined with other shardings. Metadata are to not be populated when type == TUPLE and instead metadata should be set on individual tuple elements.
Used in:
REPLICATED = 0
This sharding is replicated across all devices (implies maximal, all other fields are unused).
MAXIMAL = 1
This sharding is maximal - one device runs the entire operation.
TUPLE = 2
This sharding is a tuple - only the tuple_shardings field is valid.
OTHER = 3
None of the above; tile_shape and tile_assignment are both used.
MANUAL = 4
This op is manually sharded: the shapes are already partitioned and the partitioner should not change this op.
message
xla_data.proto:94Describes the padding configuration for Pad operation. The padding amount on both edges as well as between the elements are specified for each dimension.
Used in:
repeated dimensions = 1
The padding configuration for all dimensions.
message
xla_data.proto:96Describes the padding configuration for a dimension.
Used in:
int64 edge_padding_low = 1
Padding amount on the low-end (next to the index 0). May be negative.
int64 edge_padding_high = 2
Padding amount on the high-end (next to the highest index). May be negative.
int64 interior_padding = 3
Padding amount between the elements. May not be negative.
Used in:
PADDING_INVALID = 0
PADDING_VALID = 1
Only valid portion of the base are covered.
PADDING_SAME = 2
Extra is added to produce same output size as the input.
message
xla_data.proto:711Describes whether all data-parallelism replicas will receive the same parameter data at each buffer.
Used in:
repeated bool replicated_at_leaf_buffers = 1
A list of boolean values for the flattened leaf buffers. Each value indicates whether the corresponding leaf buffer is replicated. If this field is empty, it means no buffer is replicated. Otherwise, the number of elements in this field must match the number of leaf buffers in the HLO instruction's shape.
message
xla_data.proto:696Used to indicate the precision configuration. It has backend specific meaning.
Used in:
repeated PrecisionConfig.Precision operand_precision = 1
Used in:
DEFAULT = 0
HIGH = 1
HIGHEST = 2
Primitive types are the individual values that can be held in rectangular multidimensional arrays. A description of the rectangular multidimensional array dimensions / primitive type is given by Shape, below. LINT.IfChange
Used in: , ,
PRIMITIVE_TYPE_INVALID = 0
Invalid primitive type to serve as default.
PRED = 1
Predicates are two-state booleans.
S8 = 2
Signed integral values of fixed width.
S16 = 3
S32 = 4
S64 = 5
U8 = 6
Unsigned integral values of fixed width.
U16 = 7
U32 = 8
U64 = 9
F16 = 10
Floating-point values of fixed width. Note: if f16s are not natively supported on the device, they will be converted to f16 from f32 at arbirary points in the computation.
F32 = 11
BF16 = 16
Truncated 16 bit floating-point format. This is similar to IEEE's 16 bit floating-point format, but uses 1 bit for the sign, 8 bits for the exponent and 7 bits for the mantissa.
F64 = 12
C64 = 15
Complex values of fixed width.
Paired F32 (real, imag), as in std::complex<float>.
C128 = 18
Paired F64 (real, imag), as in std::complex<double>.
TUPLE = 13
A tuple is a polymorphic sequence; e.g. a shape that holds different sub-shapes. They are used for things like returning multiple values from a computation; e.g. a computation that returns weights and biases may have a signature that results in a tuple like (f32[784x2000], f32[2000]) If a shape proto has the tuple element type, it may not have any entries in the dimensions field.
OPAQUE_TYPE = 14
An opaque type used for passing context-specific data to a custom operation. Shapes of this primitive type will have empty dimensions and tuple_shapes fields. (OPAQUE would be a better name for this identifier, but that conflicts with a macro defined in windows.h.)
TOKEN = 17
A token type threaded between side-effecting operations. Shapes of this primitive type will have empty dimensions and tuple_shapes fields.
The type optimization profiles in use.
Used in:
INVALID = 0
WINDOW = 1
FLAG = 2
INTEGER = 3
message
xla_data.proto:234Shape of the parameters and output of a computation (like a traditional function signature).
Used in: , , ,
repeated parameters = 1
optional result = 2
repeated string parameter_names = 3
Used in:
RNG_DEFAULT = 0
Backend dependent default algorithm.
RNG_THREE_FRY = 1
RNG_PHILOX = 2
Next: 2
Used in:
RNG_INVALID = 0
RNG_UNIFORM = 1
Creates a uniform-distribution-generated random number on the semi-open interval [parameter[0], parameter[1]).
RNG_NORMAL = 2
Creates a normal-distribution-generated random number with mean parameter[0] and standard deviation parameter[1].
message
xla_data.proto:682Describes the replica groups in a cross replica op (e.g., all-reduce and all-to-all).
Used in:
repeated int64 replica_ids = 1
The ids of the replicas that belongs to the same group. The ordering of the ids matters in some ops (e.g., all-to-all).
message
xla_data.proto:510Describes the dimension numbers for a scatter operation. All the fields are similar to the corresponding fields in GatherDimensionNumbers. Differences are noted below.
Used in:
repeated int64 update_window_dims = 1
The set of dimensions in the updates shape that are window dimensions.
repeated int64 inserted_window_dims = 2
The set of window dimensions that must be inserted into the updates shape.
repeated int64 scatter_dims_to_operand_dims = 3
int64 index_vector_dim = 4
message
xla_data.proto:196A shape describes the number of dimensions in the array, the size of each dimension, and the primitive component type. Tuples are a special case in that they have rank zero and have tuple_shapes defined. See the XLA documentation for more information on shapes and layouts. LINT.IfChange
Used in: , , , , , , , , , , , , , , , , , , ,
PrimitiveType element_type = 2
The element type for this shape.
repeated int64 dimensions = 3
The size (number of elements) for each dimension, or an upper bound on the size if the dimension is dynamic. In XLA, dimensions are numbered from 0 to N-1 for an N-dimensional array. The first element of 'dimensions' is the size of dimension 0, the second element is the size of dimension 1, and so forth. Empty list indicates a scalar. If the respective element in 'is_dimension_dynamic' is true then the value in this field represents an upper bound on the size of the dimension.
repeated tuple_shapes = 4
For tuples only, the shapes of constituent shapes in the tuple sequence.
optional layout = 5
The layout used to back this shape.
repeated bool is_dynamic_dimension = 6
For arrays, this indicates whether or not each dimension is dynamically-sized. The number of elements in this repeated field should be zero (indicating that no dimensions are dynamic) or equal to the number of elements in the 'dimensions' field.
message
xla_data.proto:689Describes the source target pair in the collective permute op.
Used in:
int64 source = 1
int64 target = 2
message
xla_data.proto:126Describes 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.
Used in:
repeated int64 dimensions = 1
Number of elements in each dimension of the tile. It's ordered from the most major dimension of the tile to the most minor dimension of the tile. The dimensions correspond to a suffix of the dimensions of the shape being tiled.
message
xla_data.proto:603Used in:
bool left_side = 1
If true, solves ax = b. If false, solves xa = b.
bool lower = 2
If true, 'a' is lower triangular. If false, 'a' is upper triangular.
bool unit_diagonal = 3
If true, the diagonal elements of 'a' are assumed to be 1 and not accessed.
TriangularSolveOptions.Transpose transpose_a = 4
Should we transpose or use the adjoint of 'a'?
Used in:
TRANSPOSE_INVALID = 0
NO_TRANSPOSE = 1
Don't transpose 'a'.
TRANSPOSE = 2
Transpose 'a'.
ADJOINT = 3
Complex conjugate and transpose 'a'.
message
xla_data.proto:729A backend-config for kWhile loops that stores the loop's trip count, if it is known. This is useful for backends that can implement a `for i in 0..N` loop more efficiently than a `while` loop. For example, on GPUs, we can implement a `for i in 0..N` loop by enqueueing the kernels for the loop body N times, whereas implementing a `while` loop requires a host-device sync on each iteration.
optional known_trip_count = 1
This indirection lets us distinguish between known-trip-count == 0 and unknown-trip-count.
message
xla_data.proto:730Used in:
int64 n = 1
message
xla_data.proto:471Describes the windowing in an operation such as convolution. The window is moved across a base area and for each position of the window a computation is performed. The field below describes the window and the movement of the window across a base area.
Used in: ,
repeated dimensions = 1
message
xla_data.proto:424Used in:
int64 size = 1
The size of the window in this dimension. For a rectangle, this would be the width or height.
int64 stride = 2
The stride at which the window moves across the base area in this dimension. In other words, this is the spacing between different positions of the window in this dimension.
int64 padding_low = 3
If positive, means the amount of padding to add to the base area at the low end of this dimension; if negative, its negative means the number of elements removed from the low end of this dimension. For example, in the horizontal dimension of a rectangle, this would be the number of padding values to pad on the left, given that indices increase when going right. The actual padding value depends upon the context. Convolution pads with zeros. ReduceWindow and SelectAndScatter pads with the reduce function's init value.
int64 padding_high = 4
As padding_low, but on the high end of this dimension. For example, in the horizontal dimension of a rectangle, this would be the number of values to pad on the right, given that indices increase when going right.
int64 window_dilation = 5
Dilation factor of the sliding window in this dimension. A dilation factor of 1 means no dilation. window_dilation - 1 no-op entries ("holes") are implicitly placed between each kernel element. This value may not be less than 1. See documentation for convolution.
int64 base_dilation = 6
Dilation factor of the base area in this dimension. A dilation factor of 1 means no dilation. base_dilation - 1 no-op entries ("holes") are implicitly placed between each base area element. This value may not be less than 1. See documentation for convolution.
bool window_reversal = 7
Window reversal means that this dimension was logically reversed before the operation.