Remove usage of jax2tf.PolyShape in haiku.

jax2tf.PolyShape has been deprecated since January 2024.
Instead, we can use simple strings.

PiperOrigin-RevId: 707033476

haiku/_src/reshape_test.py

@@ -87,7 +87,7 @@ def f(inputs):
     # We convert `f` using `jax2tf` with undefined shape
     converted_f = jax2tf.convert(
         apply_fn,
-        polymorphic_shapes=[None, None, jax2tf.PolyShape("_", "T", ...)],  # pytype: disable=wrong-arg-count
+        polymorphic_shapes=[None, None, "_, T, ..."],  # pytype: disable=wrong-arg-count
         with_gradient=True,
     )