[JAX] Fix code that compares PartitionSpecs with tuples.
An upcoming change to JAX will make comparisons between PartitionSpecs and raw tuples fail. Fix code that would fail under that change by adding explicit PartitionSpec conversions.
PiperOrigin-RevId: 899076494
