Transformer Model Encoder for sequence to sequence translation.

HuggingFace tokenizer wrapper

Parameters for AutoAugment and RandAugment.

Transformer MLP / feed-forward block.

A class for logging metrics.

Transformer MLP / feed-forward block.

Resize and Keep Ratio Copy & paste from `timm`

HuggingFace tokenizer wrapper

Crops the given image at the center. If the image is torch Tensor, it is expected to have [..., H, W] shape, where ... means an arbitrary numb

Measures time and reports progress, hyper-specific to our train loops. Some concepts: 1. This differentiates between three "types" of time: -

A Transformer Encoder with interleaved self-attention and cross-attention. This architecture processes a primary sequence `x` (e.g., text) by alt

Transformer Model Encoder for sequence to sequence translation.

Custom tokenizer using WordPiece-based subword tokenization

Transformer Model Encoder for sequence to sequence translation.

HuggingFace model adapter

Subclass torch's LayerNorm (with cast back to input dtype).

Multihead Attention Pooling.

Applies Mixup and/or Cutmix to a batch of images. - Mixup: https://arxiv.org/abs/1710.09412 - Cutmix: https://arxiv.org/abs/1905.04899 Implement

A ResNet class that is similar to torchvision's but contains the following changes: - There are now 3 "stem" convolutions as opposed to 1, wi

Multi-head dot-product attention with optional TPU Flash Attention. This class extends Flax's default MHA to include support for a highly opt

https://arxiv.org/abs/2212.00794

Applies RandomErasing to a single image. Reference: https://arxiv.org/abs/1708.04896 Implementation is inspired by https://github.com/rwightman/

Sigmoid Loss for Language Image Pre-Training (SigLIP) - https://arxiv.org/abs/2303.15343 @article{zhai2023sigmoid, title={Sigmoid loss for

timm model adapter

partial(fn, **kwargs) that defines hash and eq - to help with jit caches. This is particularly common in evaluators when one has many evaluator i

ViT model.

ViT model.

Core implementation of the image-text autoregressive Transformer. Attributes: fusion_style: Method to combine image and text features.

A image-text autoregression Transformer model.

Apply Color Jitter to the PIL image with a specified probability.

Apply Gray Scale to the PIL image with a specified probability.

Adds positional embeddings to the inputs, supports caching for decode. Attributes: decode: whether to run in single-position autoregressive mod

CLS token pooling NOTE: this is equivalent to ClsPooler above with use_pooler_output=False

CLS token pooling

A Transformer block with cross-attention followed by an MLP. This block takes two inputs: `x` (typically text embeddings) which acts as the q

Single transformer encoder block (MHSA + MLP).

A custom transformer that can use different block types.

The API that any data source should implement.

Use TFDS as a data source.

Single transformer encoder block (MHSA + MLP).

Single transformer encoder block (MHSA + MLP).

Image/text retrieval evaluator.

Contrastive evaluator.

Zero-shot classification evaluator.

Image augmentation class for applying image distortions.

Decorator for preprocessing ops, which adds `inkey` and `outkey` arguments. Note: Only supports single-input single-output ops.

Subclass torch's LayerNorm to handle fp16 (by casting to float32 and back).

Max pooling

Mean pooling

A multi-modal model combining an image encoder and a text decoder. This class serves as a container that orchestrates the forward pass through

Multi-head dot-product attention. Attributes: num_heads: number of attention heads. Features (i.e. inputs_q.shape[-1]) should be di

Implements global Registry. Authors: Joan Puigcerver (jpuigcerver@), Alexander Kolesnikov (akolesnikov@)

HuggingFace tokenizer wrapper for SigLIP T5 compatible sentencepiece vocabs