Update Flax TPU tests (#3069)

Commit

3 years ago

Update Flax TPU tests (#3069) Update Flax TPU tests. Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>