Skip to content

Fix Gemma4 Array Mask Indexing#46203

Merged
vasqu merged 7 commits into
huggingface:mainfrom
petecao:petecao/fix-gemma4-indexing
May 26, 2026
Merged

Fix Gemma4 Array Mask Indexing#46203
vasqu merged 7 commits into
huggingface:mainfrom
petecao:petecao/fix-gemma4-indexing

Conversation

@petecao
Copy link
Copy Markdown
Contributor

@petecao petecao commented May 25, 2026

What does this PR do?

This PR replaces the in-place boolean indexing assignment (llm_input_ids[multimodal_mask] = ...) with an out-of-place torch.where operation in the Gemma4 modular forward pass. This fixes a jax.errors.NonConcreteBooleanIndexError: Array boolean indices must be concrete.

  • [ X ] I confirm that this is not a pure code agent PR.

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • [ x ] Did you read the contributor guideline,
    Pull Request section?
  • Was this discussed/approved via a Github issue or the forum? Please add a link
    to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?

Who can review?

@vasqu
Copy link
Copy Markdown
Contributor

vasqu commented May 26, 2026

run-slow: gemma4

Copy link
Copy Markdown
Contributor

@vasqu vasqu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks safe enough to me just sanity checking with our slow ci. Thanks!

@vasqu
Copy link
Copy Markdown
Contributor

vasqu commented May 26, 2026

GH actions are not working, will recheck later 😢

@vasqu
Copy link
Copy Markdown
Contributor

vasqu commented May 26, 2026

run-slow: gemma4

@github-actions
Copy link
Copy Markdown
Contributor

[For maintainers] Suggested jobs to run (before merge)

run-slow: gemma4

@github-actions
Copy link
Copy Markdown
Contributor

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/gemma4"]
quantizations: []

@github-actions
Copy link
Copy Markdown
Contributor

CI Results

Workflow Run ⚙️

Commit Info

Context Commit Description
RUN fd9f888d workflow commit (merge commit)
PR 9f5e3ad7 branch commit (from PR)
main 6db6fdb9 base commit (on main)

Model CI Report

1 new failed tests from this PR 😭

  • gemma4:
    tests/models/gemma4/test_modeling_gemma4.py::Gemma4IntegrationTest::test_export_text_only (❌ ⟹ ❌)

@vasqu vasqu enabled auto-merge May 26, 2026 15:52
@vasqu vasqu added this pull request to the merge queue May 26, 2026
Merged via the queue into huggingface:main with commit 5a4b70f May 26, 2026
23 of 24 checks passed
yuchenxie4645 pushed a commit to yuchenxie4645/transformers that referenced this pull request May 28, 2026
* replae boolean indexing with torch.where

* Fix consistency errors

* Resolve runtime crash

* Fix Python formatting

---------

Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>
kashif pushed a commit to kashif/transformers that referenced this pull request Jun 1, 2026
* replae boolean indexing with torch.where

* Fix consistency errors

* Resolve runtime crash

* Fix Python formatting

---------

Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants