Returns a ModelParams subclass with the given `class_key`. Args: class_key: string key of the ModelParams subclass to return. Returns:

Returns (fan_in, fan_out) of a weight variable of the give shape.

Lists devices in the job. Args: job_spec: A param object specifying a job in a training cluster. Returns: Returns a 2D np string

r"""Converts anchor_boxes and residuals to predicted bboxes. This converts predicted residuals into bboxes using the following formulae::

Return the differences between this object and another as a string. Args: other: The other Params object. Returns: A string of d

Writes to a Hyperparams proto. Serializes the Hyperparams into a proto that can be then written to disk or sent over the network. Note that s

Returns updated state from theta.

Conv2D layer.

Creates one of several checkpointer versions. Args: train_params: the training hyperparams. TODO(jlipschultz): unify with BaseRunn

Layer norm with default params.

Load a specified set of bounding boxes. If no boxes are found, return None. Note that we do *not* specify 'difficulty' as an option due to h

Add RoPE (Rotary Position Embedding) [1] to q/k proj vector. [1] https://arxiv.org/abs/2104.09864 Args: theta: A `.NestedMap` object c

Runs `loop_func`, retrying on expected errors. Args: job_name: string job name. loop_func: callable to run and retry on expected erro

Apply sequence to only update the value at a specific key. Note: This is the more general version of SeqOnKey. SeqOnKey automatically extract

Initializes SRUCell.

Shallow clone of module `m`.

Update WER when adding one pair of strings: (hypothesis, reference). Args: hypothesis: Hypothesis string. reference: Reference string

Attention with multiple attention heads. Keys, values share same dimensionality params.self.params.attention_key_value_dim. Args:

Constructs the inference graph for eval decoding. Args: input_batch: The input batch. A `NestedMap` of tensors. Or, if input batch

Converts string `text` to integer ids and the encoded string. Encoding includes prefixing the beginning-of-word token to each word. Returns:

Compute the encoding for each point. Args: unused_theta: The variables for this task. points_xyz: a floating point tf.Tensor with sha

Gets the value for nested `key`, returns `default` if key does not exist. Names with underscores will be considered as one key. Args:

Returns params for job `job_name` on the dataset `dataset_name`.

Syntactic sugar for asserting that tensor has the expected shape. Args: tensor: A Tensor. expected_shape: A Python list or a 1D tensor. Ele

r"""Scaled Huber (SmoothL1) Loss. This function wraps tf.losses.huber_loss to rescale it by 1 / delta, and uses Reduction.NONE. This sca

Add tensor 't' to 'named_tensors' at 'path'. A path may be a name or a path into a NestedMap. For instance, StoreTensor('a.b.c', [1]), is equ

Same as ToText but encodes both params and their types.

Helper for '->wq,wk,wv,wo' Graph edge.

Gets sample wav file pcm samples. Returns: (sample_rate, mono_audio) where mono_audio is of shape [batch (=1), samples].

Reset TF status. The implementation differs in Graph and Eager mode. In Graph mode, the status is renewed by entering a new graph. In Eager mode,

Resets object state. In particular, removes the reference to the TPUEmbedding object, if present, necessary for re-running unit tests with a

Adds a bbox. Args: img_id: A unique image identifier. score: The confidence score. box: [1 x 7] numpy array. difficulty:

Common input params.

Create the host-side enqueue ops. This should be called in an outer non-TPU context. Note: This is only used for the TF1 infeed based traini

Set up the input pipeline for TPU in Graph mode.

Returns layer params to compute (outputs, scalar_aux_loss).

If vars are not initialized, restore from checkpoint.

Returns Params for the training dataset.

Replaces every nested value x in 'v' with fn(x) and returns the result.

Returns a copy of this `.NestedMap` with fn applied on each key/value.

(inputs, paddings) -> (encoded, paddings). Args: name: the string name of the encoder layer params. stride: To use every k-th token,

Creates a set of empty initialized tensors. Args: nmap: A `.NestedMap` of tensors. Returns: A `.NestedMap` of tensors. Each tensor has t

Repeat each value according to num hyps. Args: source_enc_len: source encoder length; int [batch]. num_hyps_per_beam: number of hypot

Applies a FC layer to `features` key, emits a `NestedMap`.

Instantiates a model object for a particular task. MultiTaskModels can accept a shared_model parameter, but SingleTaskModels cannot, so we ha