opentau.utils.transformers_patch

Module for patching transformers

Most patches come from the branch fix/lerobot-openpi

Functions

patched_gemma_config_init(self, *args[, ...])	Initializes the GemmaConfig with added ADARMS support.
patched_gemma_decoder_layer_init(self, ...)	Initializes a GemmaDecoderLayer with potential ADARMS support.
patched_gemma_model_init(self, config)	Initializes the GemmaModel with potential ADARMS support.
patched_gemma_pretrained_model_init_weights(...)	Initializes the weights of the GemmaPreTrainedModel.
patched_paligemma_model_get_image_features(...)	Obtains image last hidden states from the vision tower and apply multimodal projection.

Classes

PatchedGemmaRMSNorm(dim[, eps, cond_dim])

RMS normalization with optional adaptive support (ADARMS).

class opentau.utils.transformers_patch.PatchedGemmaRMSNorm(dim: int, eps: float = 1e-06, cond_dim: int | None = None)

Bases: Module

RMS normalization with optional adaptive support (ADARMS).

__init__(dim: int, eps: float = 1e-06, cond_dim: int | None = None)

Initializes the PatchedGemmaRMSNorm.

Parameters:

• dim – The dimension of the input tensor.

• eps – The epsilon value for numerical stability.

• cond_dim – The dimension of the conditioning vector (if using ADARMS).

extra_repr() → str

Returns the extra representation of the module.

forward(x: Tensor, cond: Tensor | None = None) → Tuple[Tensor, Tensor | None]

Forward pass of the normalization layer.

Parameters:

• x – The input tensor.

• cond – The conditioning tensor for adaptive normalization.

Returns:

A tuple containing the normalized tensor and the gate tensor (if applicable). If cond is None, the gate tensor will be None.

Raises:

ValueError – If cond dimension does not match the configured cond_dim.

opentau.utils.transformers_patch.patched_gemma_config_init(self, *args, use_adarms: bool = False, adarms_cond_dim: int | None = None, **kwargs)

Initializes the GemmaConfig with added ADARMS support.

Parameters:

• self – The GemmaConfig instance.

• *args – Variable length argument list.

• use_adarms – Whether to use Adaptive RMS normalization.

• adarms_cond_dim – The dimension of the conditioning vector for ADARMS.

• **kwargs – Arbitrary keyword arguments.

opentau.utils.transformers_patch.patched_gemma_decoder_layer_init(self, config: GemmaConfig, layer_idx: int)

Initializes a GemmaDecoderLayer with potential ADARMS support.

Parameters:

• self – The GemmaDecoderLayer instance.

• config – The configuration object.

• layer_idx – The index of the layer.

opentau.utils.transformers_patch.patched_gemma_model_init(self, config: GemmaConfig)

Initializes the GemmaModel with potential ADARMS support.

Parameters:

• self – The GemmaModel instance.

• config – The configuration object.

opentau.utils.transformers_patch.patched_gemma_pretrained_model_init_weights(self, module: Module)

Initializes the weights of the GemmaPreTrainedModel.

Parameters:

• self – The GemmaPreTrainedModel instance.

• module – The module to initialize.

opentau.utils.transformers_patch.patched_paligemma_model_get_image_features(self, pixel_values: FloatTensor) → Tensor

Obtains image last hidden states from the vision tower and apply multimodal projection.

Two deviations from stock HuggingFace:

• Drops the / sqrt(text_config.hidden_size) scaling applied after the multi-modal projector (matching the original Pi0 fork).

• Supports native input resolutions. Stock SigLIP silently floor-crops any sub-patch remainder in its stride-patch_size conv (e.g. 180×320 with patch 14 loses 12 pixel rows and columns) and then fails on the fixed 224×224 position-embedding table. Here the pixels are padded up to the next patch multiple (bottom/right, black in the [-1, 1] input range) and interpolate_pos_encoding is enabled whenever the resulting patch grid differs from the config grid. At the config resolution (224×224) both steps are no-ops and the output is bit-identical to before. Determinism caveat: when the vision tower is trainable (freeze_vision_encoder=False), the interpolation backprops into the position-embedding table and CUDA’s bicubic-interpolate backward is non-deterministic (atomicAdd) — GPU native-resolution runs with an unfrozen tower are not bit-reproducible. The default (frozen tower) and the config-resolution path are unaffected.

Parameters:

• self – The PaliGemmaModel instance.

• pixel_values – The tensors corresponding to the input images, in [-1, 1]. Shape: (batch_size, channels, height, width).

Returns:

Image feature tensor of shape (num_images, image_length, embed_dim).