Returns the size (num of elements) of the tensor.

Updates history file iff we are recording given metric and jobname.

Tracks activations and creates variables for quantization as needed.

List of int32 indices for the classes that should be evaled.

Generates a Tensor of sinusoids with different frequencies. Uses the provided position tensor to generate positional embeddings. Refer to FPr

Compute the output shape given the input shape.

Update this metric with a newly evaluated image. Args: str_id: A string. Unique identifier of an image. result: A NestedMap with thes

Reshapes tensor x according to model_dim_reshape_segments.

r"""Adds a subplot from tensors using plot_fun to populate the subplot axes. Args: tensor_list: A list of tensors to be realized as numpy a

Configs for `MTEncoderV1`.

Wrapper for tf.while_loop that adds summary_tensors to loop vars.

Builds FProp graph. Args: theta: A NestedMap of Tensors, see base class. inputs: A Tensor of shape [batch, seqlen, dim0]. paddi

Generate a new checkpoint. May raise exceptions for failures or sanity checks. If using async mode, exceptions are raised in the following Sa

Fully connected layer, with optional batch norm.

Performs beam search based decoding. Args: encoder_outputs: the outputs of the encoder. num_hyps_per_beam_override: If set to a value

Returns a list of human-interpretable strings.

Allows subclasses control computation of gradients.

Returns true if self and other are compatible. If x and y are two compatible `.NestedMap`, `x.Pack(y.Flatten())` produces y and vice versa.

The params of this layer.

Returns a debug string for this `.NestedMap`.

Called by an eval job before evaluation. Can be used to write additional information to disk. Args: train_dir: Directory in which any

Identity. (t_1, ..., t_n) -> (t1, ..., t_n).

Applies layer normalization on the last dimension of 'x'. Args: theta: a NestedMap of layer params. gate_name: the name of the gate,

Returns the rank of the tensor.

Computes the context vector given the current query output. Note: `packed_src.source_vecs` are the vectors that are used to compute the atten

Returns a device function for placing ops within the cluster. Args: strategy: A string. Identifier for a placement strategy. By default, we

Sets multiple parameters. Dots in names indicate navigation into nested Params objects. We do not allow navigation into lists or dicts, and m

Apply identity transformation.

Padding aware max pooling layer, emits either a single tensor or map.

Compute RNN forward pass. Args: theta: A `.NestedMap` object containing weights' values of this layer and its children layers.

Rescales output of dynamic matmul (act * act). Args: output: output, corresponds to tf.matmul(act_lhs, act_rhs) Returns: Rescale

Get parameter. Dots in names indicate navigation into nested Params objects. We do not allow navigation into lists or dicts, and may ban thes

Registers an accumulator for this layer. An accumulator is used to propagate some state to a future point, where it is acted on (typically as

Wrapper for xla_sharding.split. Args: x: Tensor to annotate. split_dimension: xla_sharding.split arg. num_devices: xla_sharding.split a

Tokenize strs into vocab ids. Args: strs: A vector of strings. max_length: An int providing the max_length for strs. external_a

Quantizes activations for (act * act) matmul AQT style. This only scales, rounds and clips; resulting quantized acts would be either integer

Returns the device to use for 'device_index' for the current model split. Args: device_index: An int, the device index within 'model_split'

Create a small decoder for testing.

Gets the QDomain matching a given domain name. Args: domain: User specified domain name. Returns: The requested QDomain, the 'de

Adds a histogram summary for 'v' into the default tf graph.

Applies the gradient to the variable. Args: lr: A scalar or a callable that returns the learning rate. var_grad: A `.NestedMap` of (v

Compare golden value with real value. When running the bazel tests with FLAGS.update_goldens to be True, this function automatically updates the

Real fusion logic happens here. Works for single step or for the entire sequence. Args: theta: A NestedMap object containing weights f

Flattens 'x' by extracting tensors from nested structures to a list.

Gets the value for the nested `key`. Names with underscores will be considered as one key. Args: key: str of the form `([A-Za-

Returns the output value given the current state.

Creates a QWeight with weight_name and given shape. Args: weight_name: Unique name (within layer) for this weight. shape: Expected sh

Check if the runner should stop. Args: sess: tf.Session. step: The current GlobalStep. check_early_stop: Whether or not we want

Computes mat_g^alpha for a square symmetric matrix mat_g.

Performs one inference step on the RNN cell. If external_inputs is not None, it is added as another act input to the RNNCell. Args:

Finishes creation of the overall figure, returning the image summary.

Convenience function that returns the param for the given dataset name. Args: dataset: A python string. Typically, 'Dev', 'Test', etc.