Add a fast path for Python scalars to shaped_abstractify.

PiperOrigin-RevId: 618015741

Author: hawkinsp

jax/_src/api_util.py

@@ -593,10 +593,9 @@ def _shaped_abstractify_slow(x):
 
 # TODO(mattjj,yashkatariya): replace xla.abstractify with this, same behavior
 def shaped_abstractify(x):
-  try:
-    return _shaped_abstractify_handlers[type(x)](x)
-  except KeyError:
-    return _shaped_abstractify_slow(x)
+  handler = _shaped_abstractify_handlers.get(type(x), None)
+  return handler(x) if handler is not None else _shaped_abstractify_slow(x)
+
 _shaped_abstractify_handlers: dict[Any, Callable[[Any], core.ShapedArray]] = {}
 
 
@@ -619,6 +618,13 @@ def _np_scalar_abstractify(x: np.generic) -> ShapedArray:
 _shaped_abstractify_handlers.update((t, _np_scalar_abstractify)
                                     for t in numpy_scalar_types)
 
+def _python_scalar_abstractify(x: int | float | complex | bool) -> ShapedArray:
+  typ = type(x)
+  dtype = dtypes._scalar_type_to_dtype(typ, x)
+  return ShapedArray((), dtype, weak_type=typ in dtypes._weak_types)
+_shaped_abstractify_handlers.update((t, _python_scalar_abstractify)
+                                    for t in dtypes.python_scalar_dtypes)
+
 # This decorator exists to make it easier to monkey-patch APIs in JAX.
 # By default it does nothing, but it can be monkey-patched to do other things.
 def api_hook(fun, tag: str):

jax/_src/dtypes.py

@@ -623,10 +623,12 @@ def promote_types(a: DTypeLike, b: DTypeLike) -> DType:
   return np.dtype(_least_upper_bound(config.numpy_dtype_promotion.value, a_tp, b_tp))
 
 def is_weakly_typed(x: Any) -> bool:
+  if type(x) in _weak_types:
+    return True
   try:
     return x.aval.weak_type
   except AttributeError:
-    return type(x) in _weak_types
+    return False
 
 def is_python_scalar(x: Any) -> bool:
   try: