Remove noop spmd_mode check, correct type annotations in attention flax #170

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Draft

coolkp wants to merge 2 commits into main from code-health

Collaborator

coolkp commented May 1, 2025

Changes:

attention_flax.py
- fix type annotations
- change AttentionOp's scale param type to float, since it is passed as float to subsequent DotProductAttention calls
train_utils.py
- remove with jax.spmd_mode("allow_all"):
- its a noop now since enhanced barrier is switched on in JAX
- maxtext reference change
- jax


          Remove noop spmd_mode check, correct type annotations in attention flax

ff35c30

Signed-off-by: Kunjan patel <[email protected]>

coolkp requested a review from entrpn

May 1, 2025 18:02


          Fix type annotation mismatches in pil_utils.py and outputs.py

4005a8e

Signed-off-by: Kunjan patel <[email protected]>

coolkp force-pushed the code-health branch from 6b10706 to 4005a8e Compare

May 1, 2025 18:34

coolkp marked this pull request as draft

May 1, 2025 18:42

entrpn approved these changes

View reviewed changes

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet