package tensorflow.tpu
Get desktop application:
View/edit binary Protocol Buffers messages
This method requests the cached proto that the TPU execute op has been instructed to execute.
oneof key_oneof
string key = 1
uid_and_index = 2
CompilationCacheFetchTarget fetch_target = 3
message
tpu_compilation_cache.proto:23Response for GetTpuProgram RPC.
optional proto = 1
optional host_compute_metadata = 2
bool may_modify_variables = 3
optional compiler_metadata = 4
bool is_empty = 5
Whether the program is empty, which could be true for sharding/unsharding entries.
https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adadelta https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L933
Used in:
float rho = 1
float epsilon = 2
message
optimization_parameters.proto:60https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adagrad https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1634
Used in:
(message has no fields)
The Adam optimizer does not implement hyper-parameter update due to hardware limitations; use the dynamic learning rate feature instead, setting the learning rate to: user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t) Here, t is the current timestep. https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam https://github.com/tensorflow/tensorflow/blob/ab51450c817674c8ff08a7ae4f8ac50cdc4bed8b/tensorflow/python/training/adam.py#L32 Note that the code by default implements the lazy version of Adam (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/LazyAdamOptimizer) unless the use_non_lazy_adam parameter is set, in which case it implements the normal version of Adam that updates all parameters in the embedding table, even for entries that are not used in the current minibatch (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/AdamOptimizer). If use_non_lazy_adam is enabled, gradient accumulation is also required to be enabled in order to get correct results; a warning will be printed otherwise (which may change to an error in the future). If use_sum_inside_sqrt is set, the Adam variable update formula will be changed from m / (sqrt(v) + epsilon) to m / sqrt(v + epsilon**2); this option improves the performance of TPU training and is not expected to harm model quality.
Used in:
float beta1 = 3
float beta2 = 4
float epsilon = 5
bool use_non_lazy_adam = 8
bool use_sum_inside_sqrt = 10
message
optimization_parameters.proto:67Algorithm in http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf.
Used in:
bool update_accumulator_first = 1
Whether to use the updated or the old value of the accumulator when computing the effective learning rate. When update_accumulator_first is set to True, the updated value of the accumulator is used.
float max_var_update = 2
The max_var_update value to use. Set value to 0 (default) to disable using max_var_update to clip the gradient.
float max_accumulator = 3
The maximum value of the accumulator. Set max_accumulator to 0 (default) to disable using max_accumulator to clip the accumulator.
https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4358
Used in:
float rho = 1
float momentum = 2
float epsilon = 3
message
optimization_parameters.proto:8Used in:
optional lower = 1
-inf if not set
optional upper = 2
+inf if not set
Target type for compilation cache fetch operation.
Used in:
INVALID = 0
MAIN = 1
SHARDING = 2
UNSHARDING = 3
message
compilation_result.proto:11Describes the result of a TPU compilation.
error.Code status_code = 1
The error message, if any, returned during compilation.
string status_error_message = 2
repeated hlo_protos = 3
HLO proto.
message
optimization_parameters.proto:17Dynamic learning rate specification in the TPUEmbeddingConfiguration. The actual learning rates are provided as a scalar input list to the SendTPUEmbeddingGradients Op indexed by their tag specified through the following proto.
Used in:
int32 tag = 1
For tables where learning rates are dynamically computed and communicated to the TPU embedding program, a tag must be specified for the learning rate. The tag must be a non-negative integer. The total number of unique tags must be less than or equal to the number of tables in the TPU embedding configuration (a table does not specify any tag if it uses a constant learning rate, and specifies exactly one tag if it uses dynamic learning rates). All tags in the range [0, number_of_unique_tags) must be present in the TPU embedding configuration, i.e. a tag cannot be skipped if a different tag numerically greater than it is used in the configuration. If multiple tables specify the same tag, they *MUST* have the same dynamic learning rate, for example, their dynamic learning rate could be computed by the same TensorFlow sub-graph. The partitioning of the embedding layer would be more optimal if the number_of_unique_tags is as *LOW* as possible, i.e., if many tables share the same tag. The learning_rate input of the SendTPUEmbeddingGradients op is used to communicate dynamic learning rates to the TPU embedding program. The learning_rate input is a list of scalars where the size of the list is equal to the number of unique tags. The learning rate associated with a particular tag is specified by populating its corresponding index in the list of learning_rate scalars.
Estimator for the frequency of updates to a lookup table. It maintains an array (tf.Variable) D, where each element records the average number of global steps between two consecutive batches that hit the corresponding bucket. Once an item with bucket id i is sampled, D[i] is updated by: D[i] <- D[i] * (1 - tau) + delta[i] * tau, where tau is a learning rate between 0 and 1 (exclusive), and delta[i] = current global step - last step i is sampled. The estimated frequency (sampling rate in a batch) is thus 1 / D[i]. Elements in D are initialized with a large value max_delta. delta[i] will also be capped by this value. The exact sequence of operations used in the optimizer is shown below. last_hit_step[i] is a tf.Variable that holds the last global step at which i was sampled. delta = global_step - last_hit_step[i] clipped_delta = min(delta, params.max_delta) is_outlier = (delta >= params.outlier_threshold * D[i]) D[i] <- is_outlier ? clipped_delta : D[i] * (1 - params.tau) + clipped_delta * params.tau last_hit_step[i] <- global_step
Used in:
float tau = 1
Learning rate between (0, 1) that is used to update the array D.
float max_delta = 2
Maximum value of delta: difference between the current global step and the last global step at which the row was sampled.
float outlier_threshold = 3
Threshold used to determine whether the current update is an outlier.
float weight_exponent = 4
The weight exponent used to transform the estimated delta into weights. The transformation function is: (delta / max_delta) ^ (weight_exponent)
https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Ftrl https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/41159.pdf https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L2646 The hyperparameters for FTRL are the same as for the Keras implementation, with some additions. The "beta" parameter matches the behavior described in the second link above; "beta" / (2 * learning rate) should be added to "l2" to get equivalent behavior in the other TensorFlow implementations of this optimizer. When the multiply_linear_by_lr field is set to true, a modified formula is used for FTRL that treats the "linear" accumulator as being pre-multiplied by the learning rate (i.e., the accumulator named "linear" actually stores "linear * learning_rate"). Other than checkpoint compatibility, this is mathematically equivalent for a static learning rate; for a dynamic learning rate, it is nearly the same as long as the learning rate does not change quickly. The benefit of setting multiply_linear_by_lr to true is that the modified formula handles zero and near-zero learning rates without producing NaNs, improving flexibility for learning rate ramp-up. The allow_zero_accumulator parameter changes some internal formulas to allow zero and near-zero accumulator values at the cost of some performance; this only needs to be set if you are using an initial accumulator value of zero, which is uncommon.
Used in:
float l1 = 1
float l2 = 2
float lr_power = 3
float beta = 7
bool multiply_linear_by_lr = 6
bool allow_zero_accumulator = 8
message
tpu_compilation_cache.proto:24Used in:
bytes data = 1
Status of using gradient accumulation (doing two passes over the input gradients: one to accumulate them into a temporary array and another to apply them using the actual optimization algorithm). The extra message is to wrap the enum for scoping.
(message has no fields)
if UNSPECIFIED (default), gradient accumulation is ENABLED.
Used in:
UNSPECIFIED = 0
ENABLED = 1
DISABLED = 2
Configuration proto for hot ID optimization. This is an experimental feature that is currently disabled (by default).
Used in:
HotIdReplicationConfiguration.Status status = 1
Whether to enable or disable hot ID optimization. If UNSPECIFIED (default), hot ID optimization is DISABLED.
Used in:
UNSPECIFIED = 0
ENABLED = 1
DISABLED = 2
message
optimization_parameters.proto:48Source of learning rate to use.
Used in:
oneof learning_rate
float constant = 1
dynamic = 2
Variant of algorithm in http://proceedings.mlr.press/v44/shamir15.pdf
Used in:
float l2 = 1
float lr_power = 2
float min_servable_mdl_benefit = 3
float mdl_mix_in_margin = 4
float mdl_benefit_rampup_coeff = 5
float mdl_min_weight = 6
float benefit_revisit_scale = 7
float max_event_benefit = 8
float max_total_benefit = 9
float mdl_hard_limit = 10
bool hard_limit_min_benefit = 11
bool mdl_regularize = 12
https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L3068
Used in:
float momentum = 1
bool use_nesterov = 2
The online Yogi optimizer does not implement hyper-parameter update; use the dynamic learning rate feature instead, setting the learning rate to: user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t) Here, t is the current timestep. https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf plus some extensions based on FTRL. Note that the code by default implements the lazy version of online Yogi.
Used in:
float l1 = 1
The L1 regularization parameter (used analogously to the one in FTRL).
float l2 = 2
The L2 regularization parameter (used analogously to the one in FTRL).
float beta2 = 3
\beta_2 from Algorithm 2 in the paper.
Used in:
optional learning_rate = 13
Learning rate used for updating the embedding layer parameters.
optional clipping_limits = 2
Limits to which to clip the weight values after the backward pass; not present means no limits are applied.
optional gradient_clipping_limits = 7
Limits to which to clip the backward pass gradient before using it for updates; not present means no limits are applied.
float weight_decay_factor = 16
Amount of weight decay to apply; see weight_decay_optimizers.py for details. Almost all optimizers are supported with this option (MDL Adagrad Light does not work, and SGD does not behave as expected if it is enabled). Although there is no check, users who want weight decay will probably also want to enable gradient accumulation as well so that the decay will happen once per minibatch.
bool multiply_weight_decay_factor_by_learning_rate = 22
If true, the weight decay factor is multiplied by the current learning rate before use; this is to match the note in DecoupledWeightDecayExtension in weight_decay_optimizers.py.
GradientAccumulationStatus.Status gradient_accumulation_status = 17
Status of using gradient accumulation (doing two passes over the input gradients: one to accumulate them into a temporary array and another to apply them using the actual optimization algorithm).
optional hot_id_replication_configuration = 18
Configuration proto for hot ID replication. This is an experimental feature that is currently disabled (by default).
oneof parameters
Optimization algorithm parameters; which field is selected determines which algorithm to use.
adagrad = 3
bounded_adagrad = 19
stochastic_gradient_descent = 4
ftrl = 5
adam = 6
momentum = 8
rms_prop = 9
centered_rms_prop = 10
mdl_adagrad_light = 11
adadelta = 12
proximal_adagrad = 14
online_yogi = 20
proximal_yogi = 21
frequency_estimator = 23
user_defined_program = 24
message
dynamic_padding.proto:9A mapping between the dynamic shape dimension of an input and the arg that represents the real shape.
Used in:
int32 arg_index = 1
Input arg index with dynamic shapes.
int32 shape_index = 2
The dynamic shape dimension index.
int32 padding_arg_index = 3
The arg index that dynamic dimension maps to, which represents the value of the real shape.
https://www.tensorflow.org/api_docs/python/tf/compat/v1/train/ProximalAdagradOptimizer https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1961
Used in:
float l1 = 1
float l2 = 2
The online Yogi optimizer does not implement hyper-parameter update; use the dynamic learning rate feature instead, setting the learning rate to: user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t) Here, t is the current timestep. https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf plus some extensions based on FTRL. Note that the code by default implements the lazy version of proximal Yogi.
Used in:
float l1 = 1
The L1 regularization parameter.
float l2 = 2
The L2 regularization parameter.
float beta1 = 3
The exponential decay rate for the 1st moment estimates.
float beta2 = 4
The exponential decay rate for the 2nd moment estimates.
float epsilon = 5
A constant trading off adaptivity and noise.
https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4229
Used in:
float rho = 1
float momentum = 2
float epsilon = 3
Specification of an optimization algorithm's state variables (both the main value vector and any extra accumulators, etc.). This proto is only used internally by the TPU software and is not exposed directly to the TF model.
string name = 1
Parameter name for the state variable.
oneof usage
Usage type of this state variable.
user_defined = 2
fill_with_constant = 3
A state variable that should be filled with a constant and normally hidden from users (used for intermediate gradients being accumulated, for example).
Used in:
double initial_value = 1
A normal state variable that should be saved and restored in checkpoints and used as an input or output to non-debug TensorFlow ops.
Used in:
double padding_initial_value = 1
For padding embedding rows, this field specifies the initial value to be used. Separate initial values need to be specified for the embeddings and any extra accumulators. The initial values should be specified so as to maintain two invariants during model training: (1) The embedding vector multiplied by zero returns a vector containing all zeros. To maintain this invariant, the embedding values should never be NaNs or +-infinity. (2) Repeatedly applying the optimizer using a gradient vector of all zeros does not cause the embeddings or slot variables to become NaNs or +-infinity. The padding row is looked up when no embedding IDs are present for a feature. The semantics of embedding lookup dictate that the output must be zero under this scenario.
message
optimization_parameters.proto:82https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L629
Used in:
(message has no fields)
message
compile_metadata.proto:16This is an experimental proto used in the TF/XLA bridge to store metadata to a compile op (e.g. _TPUCompileMlir). TODO(lyandy): Deprecate proto once generic metadata proto is created.
Used in:
repeated args = 1
repeated retvals = 2
int32 num_replicas = 3
Number of replicas of the computation and number of cores in each replica. TODO(b/140721404): it may not be necessary to state the number of cores per replica here. Reconsider when replicated model-parallelism is implemented in XLA.
int32 num_cores_per_replica = 4
optional device_assignment = 8
uint64 function_library_fingerprint = 6
A fingerprint of the function library. Ensures that any functions called by the computation have matching definitions.
string session_handle = 9
Unique session identifier. Can be empty.
string guaranteed_const_fingerprint = 10
Fingerprint of guaranteed_const value. The fingerprint computation inside tpu_compile_op may be slow. The computation can be avoided by setting the fingerprint value here.
repeated padding_maps = 11
xla.DebugOptions.StepMarkerLocation step_marker_location = 12
The location of step markers that XLA compile will instrument.
int64 xla_fusion_autotuner_thresh = 13
Minimum number of batches run through the XLA graph before XLA fusion autotuner is enabled. Default value of zero disables the autotuner. The XLA fusion autotuner can improve performance by executing a heuristic search on the compiler parameters.
bool enable_automatic_model_parallelism = 14
Enables TPU compiler to add partitioning policies for inputs/outputs to the XLA computation for model parallelism.
bool use_spmd_for_xla_partitioning = 15
Whether to use XLA's SPMD or MPMD partitioner when compiler partitioning is requested.
message
compile_metadata.proto:18Description of the types and shapes of the arguments to a computation.
Used in:
DataType dtype = 1
optional shape = 2
Arg.Kind kind = 3
optional sharding = 4
The cross-core sharding of this input within each replica, e.g., assigning to one core, or replicate across all cores.
bool is_same_data_across_replicas = 5
Whether this argument will receive the same data across all replicas.
Arg.EnableXlaSharding enable_xla_sharding = 6
Whether to allow XLA to produce separate programs to shard/unshard this argument. Requires this arg to be an on-device Kind::VARIABLE, or a Kind::PARAMETER. For Kind::PARAMETER, it represents the initial value of a variable, and retval_index_for_sharding must be specified for the corresponding updated value.
int32 retval_index_for_sharding = 8
If XLA sharding is allowed on a Kind::PARAMETER, this field is used to specify the corresponding updated value in the return values. Use -1 for variables that are not updated.
bool fast_mem = 7
Whether this argument is placed on fast memory or not.
bool unrestricted_layout = 9
Whether to let XLA to decide the layout during compilation, as opposed to using a fixed layout determined by the shape.
string name = 10
Name of the node that the arg comes from.
Used in:
DISALLOWED = 0
TENTATIVE = 1
Sharding is allowed if host training loop exists.
ALLOWED = 2
Used in:
INVALID = 0
PARAMETER = 1
VARIABLE = 2
GUARANTEED_CONSTANT = 3
These are args which have been guaranteed to be constants during the session lifetime by the use of the GuaranteeConstOp (or ConstantOp).
message
compile_metadata.proto:69Description of the return values from a computation.
Used in:
optional sharding = 1
The cross-core sharding of this return value within each replica, e.g., assigning to one core, or replicate across all cores.
repeated table_descriptor = 1
TPUEmbeddingConfiguration.Mode mode = 2
int32 batch_size_per_tensor_core = 3
Number of samples in each batch of embedding layer activations sent to the TensorCore.
int32 num_hosts = 4
Number of TPU hosts used for inference/training.
int32 num_tensor_cores = 5
Number of TensorCore used for inference/training.
TPUEmbeddingConfiguration.ShardingStrategy sharding_strategy = 6
bool pipeline_execution_with_tensor_core = 7
This parameter determines if the execution of the sparse core will be pipelined with that of the TensorCore. This parameter only affects results when mode=TRAINING. If mode=INFERENCE or BACKWARD_PASS_ONLY, this parameter does not affect execution and hence, is a don't care value. false: The execution of the sparse core is not pipelined with that of the TensorCore. The forward pass of every step on the sparse core is executed only after the backward pass of the previous step is complete. And the backward pass on the sparse core is executed only after the embedding gradients have been computed on the TensorCore on every step. This ensures that the activations on every step observe the gradient updates from the previous step on both the sparse core and the TensorCore. true: The execution of the sparse core is pipelined with that of the TensorCore. The forward pass of every step on the sparse core can be executed after the forward pass of the previous step is complete without waiting for the backward pass. This improves the utilization of the sparse core allowing it to process step N+1 while the embedding gradients for step N are computed on the TensorCore. The backward pass of every step on the sparse core is executed directly after the forward pass for the next step is complete. The drawback is that embedding activations for step N+1 do not observe the embedding gradient updates from step N. This could affect model quality if step N and N+1 involve the same set of embedding IDs. However, since the embedding updates are sparse, this is generally not considered a problem.
optional output_layout = 8
Extended output layout information; deprecated and now ignored.
Mode. Should the embedding layer program be run for inference (just forward pass), training (both forward and backward pass) or just the backward_pass.
Used in:
UNSPECIFIED = 0
INFERENCE = 1
TRAINING = 2
BACKWARD_PASS_ONLY = 3
Sharding strategy of the embedding tables among the hosts. If the sharding_strategy is "mod", each id is assigned to host "id % num_hosts". For instance, 13 ids are split across 5 hosts as: [[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]]. If the sharding_strategy is "div", ids are assigned to hosts in a contiguous manner. In this case, 13 ids are split across 5 hosts as: [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]]. In both the strategies, if the id space does not evenly divide the number of hosts, each of the first "table_descriptor.vocabulary_size % num_hosts" hosts will be assigned one more id. This partitioning strategy exactly follows that in the embedding_lookup TensorFlow function at tensorflow/python/ops/embedding_ops.py.
Used in:
DIV_DEFAULT = 0
MOD = 1
Description of the various embedding tables.
Used in:
string name = 1
Name of the table.
int32 vocabulary_size = 2
Size of the vocabulary (i.e., number of rows) in the table.
int32 dimension = 3
The embedding dimension (i.e., the width of the embedding table).
int32 num_features = 4
Number of features mapped to this table.
optional optimization_parameters = 5
Details of the learning algorithm used to update the embedding parameters.
Used in:
repeated table = 1
Output locations for each feature of each table.
repeated output = 2
Shape and layout information for each tensor.
Format information for a single output tensor.
Used in:
oneof output_format
two_d = 4
Description of the output placement for one feature.
Used in:
repeated output_location = 1
Typically, only one copy of each feature is used, but multiple are allowed and the same data will be copied to all of them (with the gradients summed in the backward pass).
Location of one copy of the feature's data.
Used in:
int32 tensor_index = 1
Which output tensor this copy of the feature will go into. Must be between 0 and layout.output_size().
int32 dim0_offset = 2
Offset in dimension 0 for this feature copy. Must be between 0 and layout.output(tensor_index).dim0_size_per_sample().
int32 dim1_offset = 3
Offset in dimension 1 for this feature copy. Must be between 0 and layout.output(tensor_index).dim1_size() - table width; repeated or partially/fully overlapping values are allowed and results in the same range will be summed (with the gradients replicated in the backward pass).
Description of the output placement for features of one table.
Used in:
repeated feature = 1
Output locations for each feature loaded from this table.
Size and layout information for 2-D tensors.
Used in:
int32 dim0_size_per_sample = 2
Multiplier for output dimension 0 size; used to match legacy format that stacks features within a sample in dimension 0.
int32 dim1_size = 1
The size (in dimension 1) of this output tensor.
message
topology.proto:8Describes the geometry of a TPU mesh.
repeated int32 mesh_shape = 1
The dimensions of the TPU topology, in cores. Typically, this is a 3D topology [x, y, core], where the major dimensions correspond to TPU chips, and the minor dimension describes the number of cores on a multicore chip.
int32 num_tasks = 2
Number of TensorFlow tasks in the cluster.
int32 num_tpu_devices_per_task = 3
Number of TPU devices per task.
repeated int32 device_coordinates = 4
A flattened rank 3 int32 array with shape [num_tasks, num_tpu_devices_per_task, len(mesh_shape)]. `tasks` is the number of tasks in the TPU cluster, `devices` is the number of TPU devices per task, and the minor dimension corresponds to a position in the TPU mesh topology. Each entry [task, device, axis] gives the `axis`-th coordinate in the topology of a task/device pair.
message
tpu_compile.proto:27TPU compilation request for compiling computations into XLA HLO IR and build TPU programs.
bool use_experimental = 1
A flag reserved for using experimental version of the compilation. By default the value should be false.
bool use_mlir = 2
Use mlir to lower computation(s) to Hlo.
bool return_hlo_protos = 3
If true, returns hlo metadatas.
bool unload_cache_on_session_close = 4
If true, unloads cache on session close.
optional metadata = 5
Compilation metadata.
repeated arg_shapes = 6
Computation argument shapes.
repeated guaranteed_constants = 7
Input tensor that gives const guarantee to the TF runtime.
string mlir_module = 8
MLIR module definition.
optional fdef_lib = 9
A set of named functions used as the input to lowering to Hlo when mlir is not used.
int32 graph_def_version = 10
The version of the graph definition used to lower TF function to Hlo.
optional function = 11
Function containing the computation to compile.
Used in:
int64 uid = 1
int32 proto_index = 2
A user-defined optimizer. The contained HLO program must take the following arguments in the following order: 1. gradients 2. table weights 3. slot variables 4. an optional scalar input that is passed in via the dynamic learning rate mechanism. It must return/end in a tuple op that contains the following values in the following order: 1. new table values 2. new slot variable value The program must have shape (1,1) with dtype float32 throughout and only use HLO that operate elementwise (e.g., no reduce, no variables, no control flow and no broadcasting outside of the single scalar input). The HLO program should be written as if it were a dense update. It will be called on each row that needs an update and will applied elementwise.
Used in:
optional program = 1
repeated float padding_values = 2
Padding values for the parameter and the slots, see StateVariableSpecification.padding_initial_value below for more details on how this should be set. One value is needed for the weights and one for each slot.