Add __str__ to UnshapedArray so that whenever we print(aval), we don't see the class name by default. It only shows up when you do: repr(aval).

Weak_type in `__str__` is represented as `~int32[5, 4]` (note the tilde at the start)

PiperOrigin-RevId: 751066142

Author: yashk2810

jax/_src/api.py

@@ -358,7 +358,7 @@ def disable_jit(disable: bool = True):
   ...   return y + 3
   ...
   >>> print(f(jax.numpy.array([1, 2, 3])))  # doctest:+ELLIPSIS
-  Value of y is Traced<ShapedArray(int32[3])>with<DynamicJaxprTrace...>
+  Value of y is Traced<int32[3]>with<DynamicJaxprTrace...>
   [5 7 9]
 
   Here ``y`` has been abstracted by :py:func:`jit` to a :py:class:`ShapedArray`,

jax/_src/core.py

@@ -1688,6 +1688,9 @@ def __repr__(self):
     return '{}({}{})'.format(self.__class__.__name__, self.str_short(),
                              ", weak_type=True" if self.weak_type else "")
 
+  def __str__(self):
+    return '{}{}'.format("~" if self.weak_type else "", self.str_short())
+
   _bool    = concretization_function_error(bool)
   _int     = concretization_function_error(int, True)
   _float   = concretization_function_error(float, True)

jax/_src/tree_util.py

@@ -529,7 +529,7 @@ class Partial(functools.partial):
   >>> print_zero()
   0
   >>> call_func(print_zero)  # doctest:+ELLIPSIS
-  Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace...>
+  Traced<~int32[]>with<DynamicJaxprTrace...>
   """
 
   def __new__(klass, func, *args, **kw):

tests/core_test.py

@@ -203,6 +203,13 @@ def test_is_valid_jaxtype(self, dtype):
     else:
       self.assertFalse(core.valid_jaxtype(arr))
 
+  def test_str_aval(self):
+    aval = ShapedArray((8, 2), np.int32)
+    self.assertEqual(str(aval), "int32[8,2]")
+
+    aval = ShapedArray((8, 2), np.int32, weak_type=True)
+    self.assertEqual(str(aval), "~int32[8,2]")
+
   @parameterized.named_parameters(
       (str(i), *spec) for i, spec in enumerate(test_specs))
   def test_jit(self, f, args):

tests/lax_control_flow_test.py

@@ -1941,7 +1941,7 @@ def plus_one(p, iter_idx):
   def testScanBodyOutputError(self):
     with self.assertRaisesRegex(
         TypeError,
-        re.escape("scan body output must be a pair, got ShapedArray(float32[]).")):
+        re.escape("scan body output must be a pair, got float32[].")):
       lax.scan(lambda c, x: np.float32(0.), 0, jnp.arange(5.))
 
   def testScanMetadataError(self):