Returns a callable metric function for the model. The metrics function is for pmap-based models, where we need to normalise by doing p-sums o

Returns a callable metric function for the model. Args: split: The split for which we calculate the metrics. It should be one of th

Returns a callable metric function for the model. Args: split: The split for which we calculate the metrics. It should be one of th

Returns the set of fields that were not accessed since the last reset.

Load sharded dataset into CPU.

Helper function to assert that tensor `t` does not have `nan`, `inf`.

Helper function to assert that tensor `t` does not have `nan`, `inf`.

Call `writer` with a descriptive string and the results.

Saves predictions to CNS. Args: predictions: Serialised predictions. filename_prefix: File prefix to save the results to.

Running NMS on batched objects. Args: boxes: float arrays in shape B x n x 4. Boxes are in absolute coordinate in order of (l, t, r

Remove padding and concat and pad tokens. Removing padding tokens at the end of each caption, concat them into a single paragraph caption, an

Rescale boxes, and convert format to xywh.

Converts VideoMAE checkpoints to Jax, and checks for correctness.

Compute summary over all `datasets` that comes from config.

Saves predictions. Args: predictions: Serialised predictions. filename_prefix: File prefix to save the results to.

Text loss with label smoothing. Args: text_outputs: (text_batch_size, max_caption_length, vocab_size) matched_text: (text_batch_size,

Unwraps an image produced by wrap. Where there is a 0 in the last channel for every spatial position, the rest of the three channels in that spat

Returns 'image' with an extra channel set to all 1s.

Extracts encoder output.

Flattens a dictionary, keeping empty leaves.

Flattens a dictionary, keeping empty leaves.

Merges `restored_params` into `params`.

Add a batch of examples to the evaluator. Args: target: Target to be predicted as a Numpy array. output: Output from the model as a N

Add a batch of examples to the evaluator. Args: target: Target to be predicted as a Numpy array. output: Output from the model as a N

Apply data augmentation.

Generate caption from object features in an auto-agressive way. Args: params: pytree of network parameters. detections: dict with key

Blend image1 and image2 using 'factor'. Factor can be above 0.0. A value of 0.0 means only image1 is used. A value of 1.0 means only image2 is u

Computes the relevant metrics for all added <target, output> pairs.

Computes the relevant metrics for all added <target, output> pairs.

Returns generators for COCO object detection 2017 train & validation set. Args: batch_size: int; Determines the train batch size. eval_batc

Concert dense outputs from the network to objects. Args: outputs: dict of list of arrays. The keys should be 'heatmaps' and 'box_re

Generate fake data that resembles model `outputs` `batch` `indices`. See ObjectDetectionWithMatchingModel's loss_* functions for more details reg

Returns whether there is CLS token and the number of video frames.

Scores image features against queries. Args: image: Single uint8 Numpy image of any size. Will be converted to float and resized be

tf.identity for nested dictionary of Tensors.

Updates the train_state with data from restored_train_state.

Helper function to assert that tensor `t` does not have `nan`, `inf`.

Compute average value for each pixel over all patches containing it. For example, this can be used to compute the global boundary maps, or the

Returns the sum of classification and box regression losses. Args: logits: (batch_size, batch_size * num_max_captions, num_frames, 3).

Check if point q lies on line segment pr.

Determine the orientation of the triplet (p, q, r).

Pads the first dimension of the input `sequence` to `max_len`.

Returns a preprocessing function that operates on inputs and labels.

Removes any promptability-marker-character from a tensor of strings.

RoIAlign on multilevel features. Args: features: A sorted list of (B, Hi, Wi, C) features. boxes: (B, T, 4) boxes in XYXY format, whe

Computes linear evaluation metrics over multiple datasets. Args: repr_train_state: The train state that should be passed in as the first

Vision task with val and test splits.

Vision task with val and test splits.

Tile positional embeddings. Args: restored_posemb_grid: Positional embeddings from restored model. Shape is [n_restored_tokens, d] n_

Extract patches from an image. Args: im: Array of shape [N, C, H, W] Returns: Array of shape [N, C, R, R, H', W'] containing all

Build a tf.data.Dataset pipeline using clu.deterministic_data or DMVR. Args: split: The split to be used. start_step: Start step for GRAIN-

Build a tf.data.Dataset pipeline using clu.deterministic_data. Different from the original flexio, this function only support TFRecord. Args:

Build a tf.data.Dataset pipeline using clu.deterministic_data or DMVR. Args: split: The split to be used. start_step: Start step for GRAIN-

Convenience wrapper function for inferring the convolution output shape. Args: spatial_shape: Input (lhs) shapes for which the output shapes sh

Converts string to sequence of ints so that it is XLA compatible.

Creates a Bokeh figure showing an image.

Crops to center of image with padding then scales `image_size`.

Decodes audio spectrogram. Args: spectrogram: input mel spectrogram inflate: if True, adds a channel dimension circular_time_shift: If

Returns the stochastic depth rate for the current layer and view.

Initialization function to be jitted.

Wrapper around `make_mask_trees` that supports different input types.

With 50% prob turn the tensor negative.