jax-ml/jax · error · TypeError

input type mismatch for {_prim}

Error message

input type mismatch for {_prim}

What it means

A HiPrim application is being staged/traced with input avals that don't typematch the avals recorded when the primitive was created. The typecheck registered for call_hi_primitive_p compares each input aval with _prim.in_avals_flat and raises TypeError on mismatch.

Source

Thrown at jax/_src/hijax.py:371

  if isinstance(ct, ad_util.Zero):
    return ad_util.Zero(core.unmapped_aval(axis_data.size, d, ct.aval,
                                           axis_data.explicit_mesh_axis))
  return ct


call_hi_primitive_p = core.Primitive("call_hi_primitive")
call_hi_primitive_p.multiple_results = True
call_hi_primitive_p.skip_canonicalization = True
call_hi_primitive_p.is_high = lambda *args, _prim: True
call_hi_primitive_p.is_effectful = lambda params: bool(params['_prim'].effects)
@call_hi_primitive_p.def_effectful_abstract_eval
def _call_hi_primitive_abstract_eval(*_args, _prim):
  return _prim.out_avals_flat, _prim.effects

def _call_hi_primitive_typecheck(_ctx_factory, *in_atoms_flat, _prim):
  in_avals = [x.aval for x in in_atoms_flat]
  if not all(map(core.typematch, in_avals, _prim.in_avals_flat)):
    raise TypeError(f"input type mismatch for {_prim}")
  _prim.check()
  return _prim.out_avals_flat, _prim.effects
core.custom_typechecks[call_hi_primitive_p] = _call_hi_primitive_typecheck

def _call_hi_primitive_staging(trace, source_info, *args_flat, _prim):
  trace.frame.is_high = True
  args = tree_unflatten(_prim.in_tree, args_flat)
  ans = _prim.staging(trace, source_info, *args)
  return tree_leaves_checked(_prim.out_tree, ans)
pe.custom_staging_rules[call_hi_primitive_p] = _call_hi_primitive_staging

def _call_hi_primitive_to_lojax(*args_flat, _prim):
  args = tree_unflatten(_prim.in_tree, args_flat)
  ans = _prim.expand(*args)
  return tree_leaves_checked(_prim.out_tree, ans)
call_hi_primitive_p.to_lojax = _call_hi_primitive_to_lojax

def _call_hi_primitive_prettyprint(eqn, context, settings):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Re-create the primitive (re-trace) for the new input types instead of reusing the old instance
2. Ensure input dtypes/shapes are consistent, e.g. cast with jnp.asarray(x, dtype=...) before the call
3. Check for accidental integer-literal or weak-typed inputs (Python scalars) vs arrays

Example fix

# before
p = MyPrim(args1)  # traced with float32
p(args2_int)
# after
p = MyPrim(jnp.asarray(args2_int, args1.dtype))  # retrace with correct types

import jax.numpy as jnp
args = [jnp.asarray(a, dtype=p.in_avals_flat[i].dtype) for i, a in enumerate(args_flat)]
assert all(core.typematch(a.aval, b) for a, b in zip(traced_args, p.in_avals_flat))

def inputs_typematch(prim, args_flat) -> bool:
    import jax._src.core as core
    return all(map(core.typematch,
                   [getattr(a, 'aval', a) for a in args_flat],
                   prim.in_avals_flat))

try:
    p(*args)
except TypeError as e:
    if 'input type mismatch' in str(e):
        p = type(p)(*[jnp.asarray(a, p.in_avals_flat[i].dtype)
                      for i, a in enumerate(args)])
        return p(*args)
    raise

When it happens

Trigger: Re-binding a traced HiPrim (e.g. re-invoking a stored traced primitive, or cacheing a primitive across calls) with arguments of different dtype/shape/weak_type than the original trace inputs.

Common situations: Reusing a cached/traced primitive with a differently-typed input (e.g. int vs float, or shape change after padding), or hitting stale closures after dtype promotion.

Related errors

• The unsafe_buffer_pointer() method was called on {self._erro
• Triggering __jax_array__() during abstractification is no lo
• Cannot interpret value of type {typ} as an abstract array; i
• Mesh of an aval must be an AbstractMesh. Got {out_s.mesh} of
• Symbolic dimension '{self}' used in a context that requires

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/f57d23ef208b9280. Report an issue: GitHub.