Fix tests

[cgarciae]

What does this PR do?

Fixes a couple of tests that broke for unknown reasons.

[codecov-commenter]

Codecov Report

Merging #2253 (fda72ae) into main (481d5b8) will decrease coverage by 0.22%.
The diff coverage is n/a.

@@            Coverage Diff             @@
##             main    #2253      +/-   ##
==========================================
- Coverage   75.90%   75.67%   -0.23%     
==========================================
  Files          61       61              
  Lines        5332     5332              
==========================================
- Hits         4047     4035      -12     
- Misses       1285     1297      +12     

Impacted Files	Coverage Δ
flax/linen/transforms.py	89.68% <0.00%> (-4.77%)	⬇️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 0ebe0d9...fda72ae. Read the comment docs.