Get all tests to pass locally with no special configuration

[SamuelMarks]

Description

FIXES: #1107

Tests

python -m pytest and pytest should both now work when run from repository root.

Checklist

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

• I have performed a self-review of my code.
• 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.

---

Important: do not merge yet! - I want to solve your grain issue [ https://github.com/AI-Hypercomputer/maxtext/actions/runs/12421511705/job/34681866302 ] first; which I think this commit fixed google/grain@ce66a97...08f7cb1
EDIT: Can you ping the guys? - Looks like they just need to release anew google/grain/issues/662

[shralex]

The grain issue might be solved by #1109

[SamuelMarks]

Good to see. Also related is your recent PR #1113