r"""Class providing a method to create Gaussian orthogonal matrix. Class is responsible for constructing 2D Gaussian orthogonal arrays.

Dataset config.

Convolution with weight standardized kernel.

Attention module.

Implements a multi-layer conv block.

Two unpadded convolutions & relus. Attributes: features: Num convolutional features. padding: Type of padding: 'SAME' or 'VALID'. use_b

DETR Customized Multi-head dot-product attention. Attributes: num_heads: Number of attention heads. Features (i.e. inputs_q.shape[-1]) sh

A dummy set detection model for testing purposes.

Attention pooling layer. Attributes: features: Number of features. num_layers: Number of bottleneck blocks. stride: Stride in the Bottl

Attention pooling layer. Attributes: features: Number of features. num_layers: Number of bottleneck blocks. stride: Stride in the Bottl

FFN block for predicting bounding box coordinates.

BERT.

MultiHeadDotProductAttention that supports LoRA.

Transformer Encoder. **This is same as vit.Encoder(), but without adding positional embedding.** Attributes: num_layers: Number of layers.

Transformer encoder layer. Attributes: mlp_dim: Dimension of the mlp on top of attention block. num_heads: Number of self-attention heads.

Transformer encoder layer. Attributes: mlp_dim: Dimension of the mlp on top of attention block. num_heads: Number of self-attention heads.

Transformer encoder layer. Attributes: mlp_dim: Dimension of the mlp on top of attention block. num_heads: Number of self-attention heads.

A dummy multi-label classification model for testing purposes.

A dummy semantic segmentation model for testing purposes.

Transformer MLP / feed-forward block.

A mock action segmentation model for testing purposes.

LayoutViT Customized Multi-head dot-product attention. Attributes: num_heads: Number of attention heads. Features (i.e. inputs_q.shape[-1])

Layer for MultiScaleDeformableAttention.

Helper to go from question id (str) to id (int).

Implements a video+text single-tower backbone. Attributes: num_classes: Number of output classes. video_tower_config: The config of the vid

XViT model. Attributes: num_outputs: number of classes. mlp_dim: Dimension of the MLP on top of attention block. num_layers: Number of

Adds (optionally learned) positional embeddings to the inputs. Attributes: config: hyperparameters of the module

ResNet Stage: one or more stacked ResNet blocks. Attributes: block_size: Number of ResNet blocks to stack. nout: Number of features. fi

Bottleneck AxialResNet block. Attributes: nout: Number of output features. axial_attention_configs: Configurations of the axial attention.

Axial Self Attention module. Attributes: attention_axis: Axis of the attention axial_attention_configs: Configurations of the axial attenti

Holds beam search state data.

Holds beam search state data.

Holds beam search state data.

Holds beam search state data.

Bit ResNetV1. Attributes: num_outputs: Num output classes. If None, a dict of intermediate feature maps is returned gn_num_groups: Nu

Bottleneck layer of ResNet. Attributes: features: Number of features. stride: Stride of the down-sampled output. expansion: Expansion o

Bottleneck layer of ResNet. Attributes: features: Number of features. stride: Stride of the down-sampled output. expansion: Expansion o

Clip Transformer module. Attributes: features: Number of features. num_layers: Number of layers for each block. num_heads: Number of he

Deconvolution layer for upscaling. Attributes: features: Num convolutional features. padding: Type of padding: 'SAME' or 'VALID'. use_b

Backbone CNN for multi-scale feature extraction for DeformableDETR. Attributes: num_filters: Number of Resnet filters. num_layers: Number o

Sequence of DeformableDETRDecoderLayers. Attributes: embed_dim: Size of the hidden embedding dimension, used for query, value, embeddings

Layer of DeformableDETR decoder. Uses MultiScaleDeformableAttention for cross-attention and typical DETR dense attention for the self-attention.

Sequence of DeformableDETREncoderLayer.

Layer of DETR encoder.

DeformableDETR Transformer. Attributes: embed_dim: Size of the hidden embedding dimension. enc_embed_dim: Size of the hidden embedding dime

Finds next hidden state.

Transformer Encoder. Attributes: inputs: nd-array, Input data temporal_dims: Number of temporal dimensions in the input. mlp_dim: Dimen

Transformer Encoder modified, to use TokenLearner. Attributes: num_layers: Number of layers. mlp_dim: Dimension of the MLP on top of the at

Transformer Encoder modified, to use TokenLearner + TokenFuser. Attributes: num_layers: Number of layers. mlp_dim: Dimension of the MLP on

A dummy classification model for testing purposes.

A dummy encoder-decoder model for testing purposes.

A fake flax model.

A fake flax model.

A dummy regression model for testing purposes.

Class providing a method for fast attention via low rank decomposition. Class is responsible for providing a method <dot_product_attention> for fas

Parameters for AutoAugment and RandAugment.

Creates sinusoidal positional embeddings for inputs.

Simple input projection layer. Attributes: embed_dim: Size of the output embedding dimension. num_groups: Number channel groups for group n

Looks up class names from integer labels and adds them as text features. Attributes: tfds_name: The TFDS name of the dataset, used to determi

Audio-Visual Fusion Transformer model for Video. Attributes: mlp_dim: Dimension of the mlp on top of attention block. num_layers: Number of

Audio-Visual Fusion Transformer model for Video. Attributes: mlp_dim: Dimension of the mlp on top of attention block. num_classes: Number o

MLP with pre-norm.

Transformer MLP / feed-forward block.

Transformer MLP / feed-forward block.

Transformer MLP / feed-forward block.

Transformer MLP / feed-forward block. Attributes: config: hyperparameters of the module out_dim: optionally specify out dimension.

A mock moment retrieval model for testing purposes.

A mock temporal localization model for testing purposes.

Dataclass to keep track of state of training. Parameters are separated into frozen and learned parameters. The state of training is structured a

Dataclass to keep track of state of training. Parameters are separated into frozen and learned parameters. The state of training is structured a

Point Cloud Transformer Encoder.

Predicts deformable attention sampling offsets.

ResNet architecture. Attributes: num_outputs: Num output classes. If None, a dict of intermediate feature maps is returned. num_filte

ResNetV1 based on big_vision/models/bit.py. This variant makes the root_block optional. Attributes: num_classes: Number of output channels f

ResNet Stage: one or more stacked ResNet blocks. Attributes: block_size: Number of ResNet blocks to stack. nout: Number of features. fi

Bottleneck ResNet block. Attributes: nout: Number of output features. strides: Downsampling stride. dilation: Kernel dilation. bott

Defines a simple convolutional neural network. The model assumes the input shape is [batch, H, W, C]. Attributes: num_outputs: Number of out

ViT model for Video with factorized space-time attention.

Dropout with two paths.

Stochastic depth with two paths.

Holds beam search state data.

Holds beam search state data.

Holds beam search state data.

Text-query-based OWL-ViT model. This module computes joint text and image embeddings which are then used for localized prediction of bounding box

Convolution with tied weights. Attributes: filters: number of filters kernel_size: filter size strides: strides padding: 'SAME' or

TokenLearner module using MHA. Attributes: num_tokens: Number of tokens to generate. num_heads: Number of heads to use for the dot product

TokenLearner module. This is the module used for the experiments in the paper. Attributes: num_tokens: Number of tokens.

TokenLearner module Version 1.1, using slightly different conv. layers. Instead of using 4 conv. layers with small channels to implement spatial

Vision Transformer model with TokenLearner. Attributes: num_classes: Number of output classes. mlp_dim: Dimension of the mlp on top of at

Tokenizer for 2D inputs (e.g., images). Attributes: patches: Configuration of the patches extracted in the stem of the model. hidden_size:

Transformer module. Attributes: features: Number of features. num_layers: Number of layers for each block. num_heads: Number of heads.

Transformer module. Attributes: features: Number of features. num_layers: Number of layers for each block. num_heads: Number of heads.

Performs layer-dropout (also known as stochastic depth). Described in Huang & Sun et al, "Deep Networks with Stochastic Depth", 2016 https://ar

Performs layer-dropout (also known as stochastic depth). Described in Huang & Sun et al, "Deep Networks with Stochastic Depth", 2016 https://ar

Extends the COCO API to (optionally) support panoptic annotations.

Vision Transformer model. This differs from scenic.projects.baselines.vit in that -- Positional embeddings are added before the transformer b