flax

[JAX] Fix code that compares PartitionSpecs with tuples.

#5410

Merged