Initialization function to be jitted.

Initialization function to be jitted.

Initialization function to be jitted.

Initialization function to be jitted.

Loading json file.

A versatile vmap-like function.

Match and label boxes in each cascade stages. Args: proposals: array (B, samples_per_image, 4). gt_boxes: array (B, max_gt_boxes, 4).

Merges tokens from each view along the channel axis.

Merges tokens from each view along the time axis.

Converts labels from multi-label to multi-hot representation. In the multi-label representation, labels have shape [..., max_num_labels_per_insta

A loop implementation of multilevel ROIAlign.

Computes the numartor.

Translates panoptic annotations to standard annotations. Panoptic annotations have one extra level of nesting compared to detection annotations (

Return the function that corresponds to `name` and update `level` param.

Prepare the cost matrix. To speed up computational efficiency of the algorithm, all weights are shifted to be non-negative. Each element is reduc

Resizes the image. Args: image: Tensor; Input image. image_size: int; Image size. Returns: Resized image.

Applies the appropriate TF resizing function.

An einsum-based implementation of ROIAlign.

Run moment retrieval test steps and save results into a dict.

Run temporal localization test steps and save results into a dict.

Samples reasonable camera parameters.

Returns num_samples prompt templates uniformly at random.

Scatters the top masks back to a full-sized array of all-zero masks.

Transform original scenic arguments to optax arguments.

Fuses all video-text pairs via self attention. Args: video_tokens: A 3D float tensor of shape (batch_size, sequence_length, channel

Return the width of the shard matcher based on the number of shards.

Returns items of a dict ordered by keys.

Removes bottleneck tokens from input.

Get target labels and boxes with additional non-object appended.

Encodes the tokens in the temporal dimension.

To grid.

Creates tokens for each view. Args: x: A 5D float tensor of shape (batch, time, height, width, channels) representing the input vid

The main training loop. Args: representation_fn: A function that given a batch, returns representations that are passed to the line

Traverses nested dicts/dataclasses and emits (leaf_name, leaf_val).

Updates weights for hungarian matching using a cover. We first find the minimum uncovered weight. Then, we subtract this from all the uncovered w

Adapts PyTorch attention layers to Flax.

Add a single inference example as COCO annotation for later evaluation. Labels should not include a background/padding class, but only valid obje

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

Calculate decomposed Relative Positional Embeddings from paper:`mvitv2`. Args: attn (Tensor): attention map. q (Tensor): query q in t

Calculate decomposed Relative Positional Embeddings from paper:`MViTv2`. Args: attn (Tensor): attention map. q (Tensor): query q in t

Add a single example to the evaluator. Args: prediction: Model prediction tuple of 3 arrays: boxes, scores, classes. 'boxes' is in

Add a single example to the evaluator. Note that the postprocessing of the predictions is significantly different from the one used in DETR.

Add a single example to the evaluator. Args: task_name: Name of the task, e.g., layout. ds_name: Name of the dataset, e.g., rico.

Add a single example to the evaluator. Args: prediction: string. target: Target dictionary with keys and 'image/id'.

Add modality learned tokens.

Add modality learned tokens.

Predict segmentation mask from bounding boxes.

Adds positional embedding.

Adds positional embedding.

Add positional encodings to the input patch sequence.

Add audio spectrogram. Args: parser_builder: An instance of a builders.BaseParserBuilder. sampler_builder: An instance of a builders.Sample

Add audio spectrogram. Args: parser_builder: An instance of a builders.BaseParserBuilder. sampler_builder: An instance of a builders.Sample

Add audio spectrogram computed from waveform. Args: parser_builder: An instance of a builders.BaseParserBuilder. sampler_builder: An instan

Adds functions to process text feature to builders. This function expects the input to be either a `tf.train.SequenceExample` (with the features

Adds functions to process text feature to builders.

Adds deterministic gaussian white noise to the signal. While setting a global seeds does not guarrantee the deterministic per-sample noise genera

Aggregates gradients using JAX's pmean. Args: axis_name: Name of the axis for pmean aggregation. Returns: A `GradientTransformation`.

Gathers inputs for moment retrieval metrics from all devices.

Yields all `LocalizedNarrative` dic in a given file. Args: filename: File to load the Localized Narratives from. Yields: LN dic.

Generate object heatmaps. Args: feature_dict: dict of tensors from the dataloader num_frames: int; frames of a single model. keep_full_

Returns transformed image.

Return the transformed image. The op can consume the seed.

Applies the op to the features.

Returns transformed boxes.

Returns transformed boxes.

Apply the decoder for each modality.

Project the unmasked tokens to decoder latent space dimension.

Apply the regressor for each modality.

Apply the encoder for each modality.

Performs SpecAugment on the inputs. SpecAugment is a data augmentation technique from arXiv:1904.08779, that combines three transformations: -

Apply inference with autoregressive decoding. Apply t5x autoregressive decoding with cache using either their beam_search or temperature_samp

Autoregressive decoding text tokens.

Performs a batched gather of the data. Args: x: A [batch, num_in, ...] JTensor of data to gather from. idx: A [batch, num_out] JTensor of d

Random blur to all frames. All frames have a blur applied to them, or all do not. Args: images: `Tensor` of shape [timesteps, height, width,

Converts normalized bounding box coordinates (xyxy) to integer. Args: bbox: Tensor of shape [..., 4] image_size: Tensor or list/tuple of tw

Initializes the beam search state data structure.

Initializes the beam search state data structure.

Initializes the beam search state data structure.

Initializes the beam search state data structure.

Beam search for transformer machine translation. If `inputs` has non-zero entries, those values are not modified, i.e., the sampled values for th