flax

Migrate `transforms_test.py` from `.value` to `[...]`

#4841

Merged