SIGN IN SIGN UP

The largest collection of PyTorch image encoders / backbones. Including train, eval, inference, export scripts, and pretrained weights -- ResNet, ResNeXT, EfficientNet, NFNet, Vision Transformer (ViT), MobileNetV4, MobileNet-V3 & V2, RegNet, DPN, CSPNet, Swin Transformer, MaxViT, CoAtNet, ConvNeXt, and more

0 0 31 Python

Align gemma4_vit API with NaFlexVit per PR #2695 review

Addresses rwightman's review feedback to match the NaFlex data-pipeline
interface so existing NaFlex collators/loaders can drive Gemma4 ViT.

Public API changes on Gemma4Vit.forward / forward_features / forward_head /
forward_intermediates:

- `pixel_position_ids` -> `patch_coord` (NaFlex external (y, x) convention)
- `padding_positions`  -> `patch_valid`  (True = valid; inverse of the former)
- Accepts `Union[Tensor, Dict[str, Tensor]]` with NaFlex dict keys
  (`patches`, `patch_coord`, `patch_valid`).
- Raw `(B, C, H, W)` input unchanged; additionally supports pre-patchified
  `(B, N, P*P*C)` and `(B, N, Ph, Pw, C)` tensors from the NaFlex loader.

Internal refactor:

- `_resolve_inputs`: unwraps dict, fills default coord/valid, converts
  external (y, x) -> internal Gemma4 (x, y) once at the API boundary.
- `_naflex_to_internal_patches`: reorders NaFlex P-P-C flat layout to the
  C-P-P order that `Gemma4PatchEmbed.input_proj` expects.
- `_embed_and_encode`: shared patch_embed + RoPE + block pipeline used by
  both `forward_features` and `forward_intermediates`.
- `Gemma4PatchEmbed._project_patches`: factored out of `forward` so
  pre-patched inputs can skip the raw-image patchify step.
- `Gemma4VisionPooler` argument `pixel_position_ids` renamed to
  `position_ids` to reflect its Gemma4-internal (x, y) semantics.

Numerical equivalence (bit-perfect) preserved on the fixed-grid path:

- `gemma4_vit_e4b` pretrained: three API paths (raw tensor, raw tensor +
  explicit NaFlex coords, dict with pre-patchified P-P-C patches) produce
  identical outputs (max |delta| = 0.0).
- Mixed-resolution NaFlex batch with padding: per-sample valid-token
  outputs match solo forwards up to SDPA batch-size numerical variance.

Column-only additive attention mask retained (matches HF Gemma4 behavior
and preserves bit-perfect equivalence with reference weights).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Y
Yonghye Kwon committed
f152cb57458bb714a483bb77d01c9302dcdf3eba
Parent: a7e09c0
Committed by Ross Wightman <rwightman@users.noreply.github.com> on 4/23/2026, 5:26:34 AM