Fix test_prompt_filtering which got broken after a rebase.

[igorts-git]

Description

Fix train_rl_test that I somehow managed to break in PR #3399.

The tests passed originally, but then it broke after a rebase. I am not sure why the GitHub tests did not catch this.

Tests

Ran the affected tests:

pytest tests/post_training/unit/train_rl_test.py

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[bvandermoon]

Thanks for the fix @igorts-git. Interesting that it looks like the test did fail for the original PR: https://github.com/AI-Hypercomputer/maxtext/actions/runs/23206585089/job/67444629589?pr=3399. Maybe we need to make this test blocking on the copybara side?

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[igorts-git]

Thanks for the fix @igorts-git. Interesting that it looks like the test did fail for the original PR: https://github.com/AI-Hypercomputer/maxtext/actions/runs/23206585089/job/67444629589?pr=3399. Maybe we need to make this test blocking on the copybara side?

Weird, somehow I missed that test failure. I think it is because "pull ready" was set a few days ago, but CL presumits failed due to merge conflicts. So I rebased and the CL passed presubmits, so I clicked "submit" without noticing that the PR had test issues.