A convenience wrapper to all the loss_* functions in this class.

Generates random bounding box coordinates and turns them into a mask. Args: image: The image for which the bounding boxes are being generated.

Get non-padded:padded ratio for width/height for each mask.

Arrange loss dictionary into a metrics dictionary.

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 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 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

Fetch microbatch slice from the given batch.

Create model and train state.

Returns the number of examples in the TFRecordss.

Returns list of operation for optax optimizer.

Calculates the loss for patch consistency.

Calculates the loss for patch supervision.

Calculates the position of the rotary embedding.

Determines new labels after excluding bad classes. Excluded classes get the new label -1. Args: classes: List of tuples containing informat

Get relative positional embeddings. Args: q_size (int): size of query q. k_size (int): size of key k. rel_pos (Tensor): relativ

Get relative positional embeddings. Args: q_size (int): size of query q. k_size (int): size of key k. rel_pos (Tensor): relativ

Resizes image to a given size. Args: resize_size: either an integer H, where H is both the new height and width of the resized image, or

Get scale masks for patches with variable patchsizes.

Output (h, w) such that smallest side in image_size resizes to size.

Get inputs to the text decoder. In evaluation, we create the zero-padded text-token with the first token being BOS. In training, we handle

Get tokenizer to decode strings for eval.

Generates the initial training state.

Generates the initial training state.

Returns the visual genome train source.

Returns the visual genome train source.

Returns the VQA Accuracy for the validation / test set. Args: logits: Output of model in shape [B, L, C]. batch: Batch of data that h

Returns HN-NCE loss when considering verb hard negatives.

Returns HN-NCE loss when *not* considering verb hard negatives.

tf.identity for nested dictionary of Tensors.

Creates a split from the ImageNet dataset using TensorFlow Datasets. For the training set, we drop the last partial batch. This is fine to do bec

Updates the train_state with data from restored_train_state.

Initialize input embedding.

Initialize encoder_block_parameters.

Updates the train_state with data from `restored_train_state`.

Updates the train_state with data from `restored_train_state`. This function is writen to be used for 'finetuning' experiments. Args: t

Initialize the positional embeddings.

Initialize the train state.

Updates the train_state with data from restored_train_state. This function is writen to be used for 'fine-tuning' experiments. Here, we do some s

Initialise MBT parameters from MultiMAE with separate encoders. Args: params: PyTree of model parameters in the target model. restored_params

Updates the train_state with data from restored_train_state. This function is written to be used for 'fine-tuning' experiments. Here, we do some

Initialize the parameters from a ViT like model.

Parse an input specs into a jax.ShapeDtypeStruct.

Interpolate positional embeddings from one size to another. Args: restored_posemb_grid: Positional embeddings from restored model. Shape is

Interpolate positional embeddings from one size to another. Args: restored_posemb_grid: Positional embeddings from restored model. Shape is

Interpolate positional embeddings from one size to another. Args: restored_posemb_grid: Positional embeddings from restored model. Shape is

Intersect two line segments. Given two 2D line segments, where a line segment is defined as two 2D points. Finds the point of intersection or ret

Compute the IOU between two intervals. Args: interval_1: A tuple (start, end) containing the first interval. interval_2: A tuple (start, en

Normalizes along dimension `axis` using an L2 norm. This specialized function exists for numerical stability reasons. Args: x: An input ndarr

L2 normalize an input tensor.

Classification softmax cross entropy loss. Args: outputs: Model predictions. For the purpose of this loss, outputs must have key 'p

Classification softmax cross entropy loss and (optionally) top-1 correct. Args: outputs: Model predictions. For the purpose of this loss, o

Loads a split from the COCO dataset using TensorFlow Datasets. Args: file_patterns: the data file patterns. dataset_configs: the dataset_co

Load encoder parameters.

Loads a split from the COCO dataset using TensorFlow Datasets. Args: batch_size: int; The batch size returned by the data pipeline. train:

Logit (inverse sigmoid) function (https://en.wikipedia.org/wiki/Logit).

Returns a preprocessing function that operates on inputs and labels.

Returns a preprocessing function that operates on inputs and labels.

Preprocessing and data-augmentation functions. Currently it only supports the default data augmentation in detectron2. Args: image_set: 't

Marks a tensor of strings as not-promptable by appending a marker.

Implements a matching function. Matching function matches output detections against ground truth detections and returns indices. Args:

Project visual features if self.project_layers_name != 'none'. Args: visual_features: (batch_size, num_tokens, dim) train: bool R

Merge tensors of time and text tokens into a single tensor. Args: caption_tokens: Tensor of text tokens for each event. timestamp_token: Te

VG ground truth are overlaping. We need to merge them before evaluating. Original code: github.com/jcjohnson/densecap/blob/maste*/densecap/

Creates a split from the ModelNet40 dataset using TensorFlow Datasets. For the training set, we drop the last partial batch. This is fine to do b

Computes FAVOR+ normalizer in noncausal attention AV. Args: qs: query_prime tensor of the shape [B...,L,H,M], where M stands for the numb

Computes not-normalized FAVOR+ noncausal attention AV. Args: qs: query_prime tensor of the shape [B...,L,H,M], where M stands for the num

Sentence preprocessor.

A context manager for temporarily pausing to await arguments. Example: with chrono.paused(wait_for=some_jax_operations): # Op

Returns the pmapped train and eval steps.

Resize masks to input resolution.

Extract text features of text token.

Proprocess images. Normalize pixels for non-padded pixels.

Whitens features and adds a bias term.

Whitens features and adds a bias term.

Preprocesses a uint8 image to the format required by the model.

Used to collect predictions and targets of the whole valid/test set. Forked from scenic/projects/baselines/detr/train_utils.py Args: pred_or

Select top k elements from the given vector x, randomly breaking ties. Args: vec: input vector. k: number of elements to select prng_ke

Reload GT annotations, optionally just a subset.

Remove the cls token.

Render image using dictionary of shapes and associated RGB colors. Args: shapes: list of shape dictionaries, in order of back-object to front-o

Resize the image to min-side = size and adjust target boxes, area, mask. Args: features: dict; 'inputs' contains tf.Tensor image unbatched. 'la

Resize the image to min-side = size and adjust target boxes, area, mask. Args: features: dict; 'inputs' contains tf.Tensor image unbatched. 'la