scale * tf.random.truncated_normal(0, 1.0).

Performs scatter update: x[i] = v. A drop-in replacement for inplace_ops.alias_inplace_update ( aka tf.InplaceUpdate). Args: x: the source

Updates the metrics. Args: key: The key this (target, pred) pair belongs to. target: An array to specify the groundtruth float target

Computes SCONES loss, argmax etc. from logits.

Check the length of param_per_dataset, if it's not list, expand to list.

Asserts name's validity.

Concat x with subs(x).

Creates and returns a tf summary writer.

Returns an edge of GraphLayer to compute attenion bias for Decoders.

Returns the device to run the decoder computation.

Get the histogram proto for checkpoint write duration metric. Args: is_async: A bool indicating whether to fetch for the async write duration

A replica of Lingvo's JAX LN.

Overriding with bias-less layer norm.

Primer normalization.

Project layer. Simply X*W+b.

Main helper method.

True LN normalization.

Returns indicator vector if preconditioner exists for each axis.

Computes decoder from params and computes loss with random inputs.

Wraps `self.train_op` with an op updating exponential moving average.

[1]. Sec 2.2.

Computes gating for Mixture-of-Experts. See Top2GatingOnLogits for more details. Note that for local_dispatch original batch BLM is reshaped int

Computes the overall step rate.

Creates a dict of decoder metrics for `PostProcessDecodeOut` to update. Returns a dict mapping from string keys to `.BaseMetric` objects.

Create optimizer slot variable and add it to the given variable dict.

DecSelfAttention with relative Attention Bias. Note that attention bias (see _DecNotVisible) ensures that current position (~row) is less tha

Clean EncoderLayerStack with minimal layer nesting. E.g:: encoder/ layer_000/ ln/w/ scale self_att

Extend prefix as represented by `prefix_states` by one more step. This function is expected to be called during fast decoding of Transformer

Computes the loss between the predictions and labels.

Align image features to point cloud features. Args: theta: A `.NestedMap` object containing variable values of this task. image_featu

Find the given result from a list of expected results. Args: result: A MassOutput tuple, from running ops.mass(). expected_results: A list

Returns function corresponding to GLU variants.

Generate a random seed from a name string. Args: name: A string. Returns: An integer seed in the range [0, 2**31 - 1).

Return the global_step.

Gets an input batch.

Returns a dict of variables from the model and all its children.

Initializes a grapheme tokenizer.

Initializes a Word Piece Tokenizer.

Computes lengths of each sequence in a batch, ignoring trailing padding. Note the following isn't guaranteed due to leading paddings. PaddingsFro

Maximum empirically observed Euclidean distance in world coordinates.

Number of logarithmically space bins for examining points.

Params builder for APMetrics.

Receives a tensor from the channel.

Sends a tensor through the channel.

Replaced dropout layers in transformer with deterministic ones.

Runs any necessary cleanup (potentially blocking).

Wait for any outstanding async operations to finish.

Create TPU dequeue ops. This should only be called within a TPU context. Returns: - A NestedMap of the input batch.