Make literal's dtype print with an empty shape so that it's consistent. So 1.0:f32 -> 1.0:f32[]

PiperOrigin-RevId: 753640777

Author: yashk2810

docs/aot.md

@@ -50,7 +50,7 @@ some other features along the way. An example:
 >>> # Print the specialized, staged-out representation (as Jaxpr IR)
 >>> print(traced.jaxpr)
 { lambda ; a:i32[] b:i32[]. let
-    c:i32[] = mul 2:i32 a
+    c:i32[] = mul 2:i32[] a
     d:i32[] = add c b
   in (d,) }
 

jax/_src/api.py

@@ -2396,7 +2396,7 @@ def make_jaxpr(
       c:f32[] = sin a
       _:f32[] = sin b
       d:f32[] = cos b
-      e:f32[] = mul 1.0:f32 d
+      e:f32[] = mul 1.0:f32[] d
       f:f32[] = neg e
       g:f32[] = mul f c
     in (g,) }

jax/_src/core.py

@@ -475,7 +475,7 @@ def pretty_print(self, context: JaxprPpContext, *, print_dtype: bool = True):
     del context  # unused
     dtype = getattr(self.aval, 'dtype', None)
     if print_dtype and dtype:
-      return f'{self.val}:{dtypes.short_dtype_name(dtype)}'
+      return f'{self.val}:{self.aval.str_short(short_dtypes=True)}'
     else:
       return f'{self.val}'
 

tests/api_test.py

@@ -6679,7 +6679,7 @@ def fun(x):
       return (x, 1., np.zeros(1, dtype=jnp.float32))
 
     dtype = "f64" if config.enable_x64.value else "f32"
-    expected = f"{{ lambda a:f32[1]; b:f32[]. let  in (b, 1.0:{dtype}, a) }}"
+    expected = f"{{ lambda a:f32[1]; b:f32[]. let  in (b, 1.0:{dtype}[], a) }}"
     jaxpr = api.make_jaxpr(fun)(jnp.float32(0.))
     self.assertMultiLineStrippedEqual(expected, str(jaxpr))
 
@@ -6691,9 +6691,9 @@ def f(x):
                       x + 2.,
                       lambda xf: xf - x)
     expected = """{ lambda ; a:f32[]. let
-    b:bool[] = ge a 0.0:f32
-    c:f32[] = add a 1.0:f32
-    d:f32[] = add a 2.0:f32
+    b:bool[] = ge a 0.0:f32[]
+    c:f32[] = add a 1.0:f32[]
+    d:f32[] = add a 2.0:f32[]
     e:i32[] = convert_element_type[new_dtype=int32 weak_type=False] b
     f:f32[] = cond[
       branches=(

tests/pjit_test.py

@@ -1274,7 +1274,7 @@ def test_pretty_print_with_constant_pjit_arg(self):
                 b:f32[1] = pjit[
                   name=<lambda>
                   jaxpr={ lambda ; a:f32[1] c:f32[]. let b:f32[1] = mul a c in (b,) }
-                ] a 1.0:f32
+                ] a 1.0:f32[]
               in (b,) }
         """).strip(),
     )
@@ -1308,7 +1308,7 @@ def test_pretty_print_with_literal_outvar(self):
             { lambda ; a:f32[1]. let
                 b:i32[] c:f32[1] = pjit[
                   name=<lambda>
-                  jaxpr={ lambda ; a:f32[1]. let  in (2:i32, a) }
+                  jaxpr={ lambda ; a:f32[1]. let  in (2:i32[], a) }
                 ] a
               in (b, c) }
         """).strip(),

tests/state_test.py

@@ -361,7 +361,7 @@ def body(x_ref):
       return []
     jaxpr, _ , _, () = pe.trace_to_jaxpr_dynamic(
         wrap_init(body, 1), [shaped_array_ref((), jnp.int32)])
-    self.assertIn("a[] <- 2:i32", jaxpr.pretty_print(use_color=False))
+    self.assertIn("a[] <- 2:i32[]", jaxpr.pretty_print(use_color=False))
 
     def body(x_ref, val):
       x_ref[:, 0] = val
@@ -377,7 +377,7 @@ def body(x_ref):
       return [x]
     jaxpr, _ , _, () = pe.trace_to_jaxpr_dynamic(
         wrap_init(body, 1), [shaped_array_ref((), jnp.int32)])
-    self.assertIn("b:i32[], a[] <- a[], 2:i32", jaxpr.pretty_print(use_color=False))
+    self.assertIn("b:i32[], a[] <- a[], 2:i32[]", jaxpr.pretty_print(use_color=False))
 
     def body(x_ref, val):
       x = ref_swap(x_ref, (slice(None), 0), val)