Speed up check_jaxpr().

(check_jaxpr() is only used when debugging.)

Don't eagerly pretty print jaxprs: only do so if we are going to raise an error.
Don't eagerly form error messages. Delete typecheck_assert.

PiperOrigin-RevId: 422594126

jax/_src/lax/control_flow.py

@@ -121,12 +121,13 @@ def _abstractify(x):
   return raise_to_shaped(core.get_aval(x))
 
 def _typecheck_param(prim, param, name, msg_required, pred):
-  msg = (f'invalid {prim} param {name} of type {type(param).__name__}, '
-         f'{msg_required} required:')
-  param_str = str(param)
-  sep = os.linesep if os.linesep in param_str else ' '
-  msg = sep.join([msg, param_str])
-  core.typecheck_assert(pred, msg)
+  if not pred:
+    msg = (f'invalid {prim} param {name} of type {type(param).__name__}, '
+           f'{msg_required} required:')
+    param_str = str(param)
+    sep = os.linesep if os.linesep in param_str else ' '
+    msg = sep.join([msg, param_str])
+    raise core.JaxprTypeError(msg)
 
 
 ### fori_loop and while_loop
@@ -1340,47 +1341,45 @@ def _cond_typecheck(*avals, branches, linear):
   tc(linear, 'linear', 'tuple of bool',
      type(linear) is tuple and all(type(x) is bool for x in linear))
 
-  core.typecheck_assert(
-      len(branches) > 0,
-      'cond requires at least one branch function')
-  core.typecheck_assert(
-      len(linear) + 1 == len(avals),
-      f'cond given {len(linear)} linear flags for '
-      f'{len(avals) - 1} non-predicate operands')
+  if len(branches) == 0:
+    raise core.JaxprTypeError('cond requires at least one branch function')
+  if len(linear) + 1 != len(avals):
+    raise core.JaxprTypeError(f'cond given {len(linear)} linear flags for '
+                              f'{len(avals) - 1} non-predicate operands')
 
   jaxpr0 = branches[0]
   jaxpr0_in_avals_str = _avals_short(jaxpr0.in_avals)
   jaxpr0_out_avals_str = _avals_short(jaxpr0.out_avals)
 
   for i, jaxpr in enumerate(branches[1:]):
-    core.typecheck_assert(
-        len(jaxpr0.in_avals) == len(jaxpr.in_avals),
+    if len(jaxpr0.in_avals) != len(jaxpr.in_avals):
+      raise core.JaxprTypeError(
         f'cond branch 0 takes {len(jaxpr0.in_avals)} inputs, '
         f'branch {i+1} takes {len(jaxpr.in_avals)}')
-    core.typecheck_assert(
-        len(jaxpr0.out_avals) == len(jaxpr.out_avals),
+    if len(jaxpr0.out_avals) != len(jaxpr.out_avals):
+      raise core.JaxprTypeError(
         f'cond branch 0 outputs {len(jaxpr0.out_avals)} values, '
         f'branch {i+1} outputs {len(jaxpr.out_avals)}')
-    core.typecheck_assert(
-        all(_map(core.typematch, jaxpr0.in_avals, jaxpr.in_avals)),
+    if not all(_map(core.typematch, jaxpr0.in_avals, jaxpr.in_avals)):
+      raise core.JaxprTypeError(
         f'cond branches 0 and {i+1} have mismatching input types: '
         f'{jaxpr0_in_avals_str} vs {_avals_short(jaxpr.in_avals)}')
-    core.typecheck_assert(
-        all(_map(core.typematch, jaxpr0.out_avals, jaxpr.out_avals)),
+    if not all(_map(core.typematch, jaxpr0.out_avals, jaxpr.out_avals)):
+      raise core.JaxprTypeError(
         f'cond branches 0 and {i+1} have mismatching output types: '
         f'{jaxpr0_out_avals_str} vs {_avals_short(jaxpr.out_avals)}')
 
-  core.typecheck_assert(
-      len(avals) == 1 + len(jaxpr0.in_avals),
+  if len(avals) != 1 + len(jaxpr0.in_avals):
+    raise core.JaxprTypeError(
       f'cond called with {len(avals) - 1} non-predicate operands, '
       f'but branches take {len(jaxpr0.in_avals)} inputs')
 
   index_aval, *op_avals = avals
-  core.typecheck_assert(
-      index_aval.dtype == np.int32,
+  if index_aval.dtype != np.int32:
+    raise core.JaxprTypeError(
       f'cond called with index of type {index_aval.dtype} instead of int32')
-  core.typecheck_assert(
-      all(_map(core.typecompat, jaxpr0.in_avals, op_avals)),
+  if not all(_map(core.typecompat, jaxpr0.in_avals, op_avals)):
+    raise core.JaxprTypeError(
       f'cond branches take input types {jaxpr0_in_avals_str}, '
       f'called with operands of type {_avals_short(op_avals)}')
 
@@ -2177,8 +2176,8 @@ def _scan_typecheck(bind_time, *avals, reverse, length, num_consts, num_carry,
   tc(length, 'length', 'non-negative int',
      type(length) in length_types and length >= 0)
 
-  core.typecheck_assert(
-      len(linear) == len(avals),
+  if len(linear) != len(avals):
+    raise core.JaxprTypeError(
       f'scan param linear has length {len(linear)} for {len(avals)} operands')
 
   const_avals, init_avals, x_avals = split_list(avals, [num_consts, num_carry])
@@ -2187,20 +2186,20 @@ def _scan_typecheck(bind_time, *avals, reverse, length, num_consts, num_carry,
   carry_avals_jaxpr, _ = split_list(jaxpr.out_avals, [num_carry])
   x_avals_mapped = _map(partial(core.mapped_aval, length, 0), x_avals)
 
-  core.typecheck_assert(
-      all(_map(core.typematch, init_avals_jaxpr, carry_avals_jaxpr)),
+  if not all(_map(core.typematch, init_avals_jaxpr, carry_avals_jaxpr)):
+    raise core.JaxprTypeError(
       f'scan input carry input and output types mismatch: '
       f'\n{_avals_short(init_avals_jaxpr)}\nvs\n{_avals_short(carry_avals_jaxpr)}')
-  core.typecheck_assert(
-      all(_map(core.typecompat, const_avals_jaxpr, const_avals)),
+  if not all(_map(core.typecompat, const_avals_jaxpr, const_avals)):
+    raise core.JaxprTypeError(
       f'scan jaxpr takes input const types\n{_avals_short(const_avals_jaxpr)},\n'
       f'called with consts of type\n{_avals_short(const_avals)}')
-  core.typecheck_assert(
-      all(_map(core.typecompat, init_avals_jaxpr, init_avals)),
+  if not all(_map(core.typecompat, init_avals_jaxpr, init_avals)):
+    raise core.JaxprTypeError(
       f'scan jaxpr takes input carry types\n{_avals_short(init_avals_jaxpr)},\n'
       f'called with initial carry of type\n{_avals_short(init_avals)}')
-  core.typecheck_assert(
-      all(_map(core.typecompat, x_avals_jaxpr, x_avals_mapped)),
+  if not all(_map(core.typecompat, x_avals_jaxpr, x_avals_mapped)):
+    raise core.JaxprTypeError(
       f'scan jaxpr takes input sequence types\n{_avals_short(x_avals_jaxpr)},\n'
       f'called with sequence of type\n{_avals_short(x_avals)}')
 

jax/core.py

@@ -16,6 +16,7 @@
 import collections
 from collections import namedtuple
 from contextlib import contextmanager
+import functools
 from functools import partial, partialmethod, total_ordering
 import gc
 import itertools as it
@@ -2050,10 +2051,6 @@ def typematch(aval1: AbstractValue, aval2: AbstractValue) -> bool:
 
 class JaxprTypeError(TypeError): pass
 
-def typecheck_assert(pred, msg):
-  if not pred:
-    raise JaxprTypeError(msg)
-
 custom_typechecks: Dict[Primitive, Callable] = {}
 
 def check_jaxpr(jaxpr: Jaxpr):
@@ -2067,13 +2064,17 @@ def check_jaxpr(jaxpr: Jaxpr):
   Raises `JaxprTypeError` if `jaxpr` is determined invalid. Returns `None`
   otherwise.
   """
-  ctx = JaxprPpContext()
-  try: pp_jaxpr(jaxpr, ctx)  # side-effect on ctx, build variable names
-  except: pass
+  @functools.lru_cache(maxsize=None)
+  def ctx_factory():
+    ctx = JaxprPpContext()
+    try: pp_jaxpr(jaxpr, ctx)  # side-effect on ctx, build variable names
+    except: pass
+    return ctx
 
   try:
-    _check_jaxpr(ctx, jaxpr, [v.aval for v in jaxpr.invars])
+    _check_jaxpr(ctx_factory, jaxpr, [v.aval for v in jaxpr.invars])
   except JaxprTypeError as e:
+    ctx = ctx_factory()
     if len(e.args) == 2:
       msg, eqnidx = e.args
       jaxpr_str = str(pp_jaxpr_eqn_range(jaxpr, eqnidx - 10, eqnidx + 10, ctx))
@@ -2083,22 +2084,28 @@ def check_jaxpr(jaxpr: Jaxpr):
     msg = "\n\n".join([msg, "while checking jaxpr:", jaxpr_str])
     raise JaxprTypeError(msg) from None
 
-def _check_jaxpr(ctx: 'JaxprPpContext', jaxpr: Jaxpr,
+def _check_jaxpr(ctx_factory: Callable[[], 'JaxprPpContext'], jaxpr: Jaxpr,
                  in_avals: Sequence[AbstractValue]) -> None:
 
   def read(v: Atom) -> AbstractValue:
     if isinstance(v, Literal):
       return raise_to_shaped(get_aval(v.val))
     else:
-      typecheck_assert(v in env, f"Variable '{pp_var(v, ctx)}' not defined")
+      if v not in env:
+        ctx = ctx_factory()
+        raise JaxprTypeError(f"Variable '{pp_var(v, ctx)}' not defined")
       return env[v]
 
   def write(v: Var, a: AbstractValue) -> None:
-    typecheck_assert(v not in env, f"Variable '{pp_var(v, ctx)}' already bound")
+    if v in env:
+      ctx = ctx_factory()
+      raise JaxprTypeError(f"Variable '{pp_var(v, ctx)}' already bound")
     if not isinstance(v, DropVar):
-      typecheck_assert(typecompat(v.aval, a),
-                       f"Variable '{pp_var(v, ctx)}' inconsistently typed as "
-                       f"{pp_aval(a, ctx)}, bound as {pp_aval(v.aval, ctx)}")
+      if not typecompat(v.aval, a):
+        ctx = ctx_factory()
+        raise JaxprTypeError(
+            f"Variable '{pp_var(v, ctx)}' inconsistently typed as "
+            f"{pp_aval(a, ctx)}, bound as {pp_aval(v.aval, ctx)}")
       env[v] = a
 
   env : Dict[Var, AbstractValue] = {}
@@ -2111,20 +2118,21 @@ def write(v: Var, a: AbstractValue) -> None:
     prim = eqn.primitive
     try:
       in_avals = map(read, eqn.invars)
-      typecheck_assert(all(not isinstance(ina, ConcreteArray) for ina in in_avals),
-                       "Equation given ConcreteArray type inputs")
+      if any(isinstance(ina, ConcreteArray) for ina in in_avals):
+        raise JaxprTypeError("Equation given ConcreteArray type inputs")
       if prim in custom_typechecks:
         out_avals = custom_typechecks[prim](*in_avals, **eqn.params)
         if out_avals is None:
           out_avals = [v.aval for v in eqn.outvars]
       elif prim.call_primitive:
-        out_avals = check_call(ctx, prim, in_avals, eqn.params)
+        out_avals = check_call(ctx_factory, prim, in_avals, eqn.params)
       elif prim.map_primitive:
-        out_avals = check_map(ctx, prim, in_avals, eqn.params)
+        out_avals = check_map(ctx_factory, prim, in_avals, eqn.params)
       else:
         out_avals = check_eqn(prim, in_avals, eqn.params)
       map(write, eqn.outvars, out_avals)
     except JaxprTypeError as e:
+      ctx = ctx_factory()
       msg, = e.args
       src = source_info_util.summarize(eqn.source_info)
       msg = "\n\n".join([msg, "in equation:", str(pp.nest(2, pp_eqn(eqn, ctx))),
@@ -2142,57 +2150,58 @@ def check_eqn(prim, in_avals, params):
     out_avals = [out_avals]
   return out_avals
 
-def check_call(ctx, prim, in_avals, params):
-  typecheck_assert("call_jaxpr" in params,
-                   f"Call primitive {prim} missing 'call_jaxpr' parameter")
+def check_call(ctx_factory, prim, in_avals, params):
+  if "call_jaxpr" not in params:
+    raise JaxprTypeError(
+        f"Call primitive {prim} missing 'call_jaxpr' parameter")
   call_jaxpr = params["call_jaxpr"]
 
   # These checks also happen in recursive call, but give better errors here.
-  typecheck_assert(len(in_avals) == len(call_jaxpr.invars),
-                   f"Call primitive {prim} with {len(call_jaxpr.invars)} "
-                   f"operands cannot call jaxpr with {len(call_jaxpr.invars)} "
-                   f"inputs")
+  if len(in_avals) != len(call_jaxpr.invars):
+    raise JaxprTypeError(f"Call primitive {prim} with {len(call_jaxpr.invars)} "
+                         f"operands cannot call jaxpr with {len(call_jaxpr.invars)} "
+                         f"inputs")
   binder_avals = [v.aval for v in call_jaxpr.invars]
   for binder_aval, in_aval in zip(binder_avals, in_avals):
-    typecheck_assert(typecompat(binder_aval, in_aval),
-                     f"Call primitive {prim} passes operand {in_aval} "
-                     f"to jaxpr expecting {binder_aval}")
+    if not typecompat(binder_aval, in_aval):
+      raise JaxprTypeError(f"Call primitive {prim} passes operand {in_aval} "
+                           f"to jaxpr expecting {binder_aval}")
 
-  _check_jaxpr(ctx, call_jaxpr, in_avals)
+  _check_jaxpr(ctx_factory, call_jaxpr, in_avals)
 
   out_avals = [v.aval for v in call_jaxpr.outvars]
   return out_avals
 
-def check_map(ctx, prim, in_avals, params):
-  typecheck_assert("call_jaxpr" in params,
-                   f"Map primitive {prim} missing 'call_jaxpr' parameter")
+def check_map(ctx_factory, prim, in_avals, params):
+  if "call_jaxpr" not in params:
+    raise JaxprTypeError(f"Map primitive {prim} missing 'call_jaxpr' parameter")
   call_jaxpr = params["call_jaxpr"]
-  typecheck_assert("axis_size" in params,
-                   f"Map primitive {prim} missing 'axis_size' parameter")
+  if "axis_size" not in params:
+    raise JaxprTypeError(f"Map primitive {prim} missing 'axis_size' parameter")
   axis_size = params["axis_size"]
-  typecheck_assert("axis_name" in params,
-                   f"Map primitive {prim} missing 'axis_name' parameter")
+  if "axis_name" not in params:
+    raise JaxprTypeError(f"Map primitive {prim} missing 'axis_name' parameter")
   axis_name = params["axis_name"]
-  typecheck_assert("in_axes" in params,
-                   f"Map primitive {prim} missing 'in_axes' parameter")
+  if "in_axes" not in params:
+    raise JaxprTypeError(f"Map primitive {prim} missing 'in_axes' parameter")
   in_axes = params["in_axes"]
-  typecheck_assert("out_axes" in params,
-                   f"Map primitive {prim} missing 'out_axes' parameter")
+  if "out_axes" not in params:
+    raise JaxprTypeError(f"Map primitive {prim} missing 'out_axes' parameter")
   out_axes = params["out_axes"]
 
   binder_avals = [unmapped_aval(axis_size, axis_name, in_axis, v.aval)
                   if in_axis is not None else v.aval
                   for v, in_axis in zip(call_jaxpr.invars, in_axes)]
   for binder_aval, in_aval in zip(binder_avals, in_avals):
-    typecheck_assert(typecompat(binder_aval, in_aval),
-                     f"Call primitive {prim} passes operand {in_aval} "
-                     f"to jaxpr expecting {binder_aval}")
+    if not typecompat(binder_aval, in_aval):
+      raise JaxprTypeError(f"Call primitive {prim} passes operand {in_aval} "
+                           f"to jaxpr expecting {binder_aval}")
 
   mapped_avals = [mapped_aval(axis_size, in_axis, aval)
                   if in_axis is not None else aval
                   for aval, in_axis in zip(in_avals, in_axes)]
   with extend_axis_env(params['axis_name'], axis_size, None):
-    _check_jaxpr(ctx, call_jaxpr, mapped_avals)
+    _check_jaxpr(ctx_factory, call_jaxpr, mapped_avals)
 
   mapped_out_avals = [v.aval for v in call_jaxpr.outvars]
   out_avals = [unmapped_aval(axis_size, axis_name, out_axis, aval) if out_axis is not None else aval

jax/experimental/maps.py

@@ -951,14 +951,15 @@ def _typecheck_xmap(
   binder_in_avals = [_insert_aval_axes(v.aval, a_in_axes, local_axis_sizes)
                      for v, a_in_axes in zip(call_jaxpr.invars, in_axes)]
   for binder_in_aval, in_aval in zip(binder_in_avals, in_avals):
-    core.typecheck_assert(
-        core.typecompat(binder_in_aval, in_aval),
+    if not core.typecompat(binder_in_aval, in_aval):
+      raise core.JaxprTypeError(
         f"xmap passes operand {in_aval} to jaxpr expecting {binder_in_aval}")
 
   mapped_in_avals = [_delete_aval_axes(a, a_in_axes, global_axis_sizes)
                      for a, a_in_axes in zip(in_avals, in_axes)]
   with core.extend_axis_env_nd(global_axis_sizes.items()):
-    core._check_jaxpr(core.JaxprPpContext(), call_jaxpr, mapped_in_avals)
+    core._check_jaxpr(lambda: core.JaxprPpContext(), call_jaxpr,
+                      mapped_in_avals)
 
   mapped_out_avals = [v.aval for v in call_jaxpr.outvars]
   out_avals = [_insert_aval_axes(a, a_out_axes, local_axis_sizes)