{"record":{"id":"b87e2179e6ef89f4","repo":"jax-ml/jax","slug":"name-argument-type-error-lhs-dtype-rhs-dtyp","errorCode":null,"errorMessage":"{name} argument type error: {lhs.dtype}, {rhs.dtype}","messagePattern":"(.+?) argument type error: (.+?), (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5879,"sourceCode":"  del dimension_numbers  # unused\n  # We're mostly matching XLA's logic here, namely in shape_inference.cc and\n  # primitive_util.h's HigherPrecisionType, e.g.\n  # https://github.com/openxla/xla/blob/ea3a841768d0dcf192e5820c9b25c34c73f2226a/xla/primitive_util.h#L329\n  def type_properties(dt):\n    c = _real_dtype(dt) if dtypes.issubdtype(dt, np.complexfloating) else dt\n    return (dtypes.issubdtype(dt, np.complexfloating),\n            dtypes.finfo(c).maxexp if dtypes.issubdtype(c, np.floating) else -1,\n            dtypes.finfo(c).nmant  if dtypes.issubdtype(c, np.floating) else -1,\n            _bit_width(c),\n            not dtypes.issubdtype(c, np.unsignedinteger))\n  lhs_prop, rhs_prop = type_properties(lhs.dtype), type_properties(rhs.dtype)\n  if lhs_prop > rhs_prop:\n    result_dtype = lhs.dtype\n  elif rhs_prop > lhs_prop:\n    result_dtype = rhs.dtype\n  else:\n    if lhs.dtype != rhs.dtype:\n      raise TypeError(f'{name} argument type error: {lhs.dtype}, {rhs.dtype}')\n    result_dtype = lhs.dtype\n  has_algorithm = isinstance(precision, (DotAlgorithm, DotAlgorithmPreset))\n  return _maybe_upcast(result_dtype, preferred_element_type,\n                       check_bit_width=not has_algorithm)\n\ndef _bit_width(d):\n  if dtypes.issubdtype(d, np.inexact): return dtypes.finfo(d).bits\n  elif dtypes.issubdtype(d, np.integer): return dtypes.iinfo(d).bits\n  elif d == np.dtype('bool'): return 1\n  else: assert False, d  # should be unreachable, open an issue!\n\ndef _maybe_upcast(result_dtype, preferred_element_type, check_bit_width):\n  # replicates the logic in shape_inference.cc's MaybeUpcast\n  if (preferred_element_type is None or\n      result_dtype == preferred_element_type):\n    return result_dtype\n  if (check_bit_width and not dtypes.issubdtype(result_dtype, np.floating) and\n      _bit_width(preferred_element_type) < _bit_width(result_dtype)):","sourceCodeStart":5861,"sourceCodeEnd":5897,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5861-L5897","documentation":"Raised in the dot_general dtype rule when lhs and rhs have equal 'type properties' (precision class) but different dtypes, e.g. two different-width floats or mismatched signed/unsigned ints of the same rank. JAX refuses to silently pick a winner between same-class types.","triggerScenarios":"jax.lax.dot_general (or matmul paths hitting it) with operands like f32 @ f64, i32 @ u32, or bf16 @ f16 — same precision class, different dtypes.","commonSituations":"Mixing a model stored in float64 with float32 inputs; mixed-precision attention where q/k dtypes differ; implicit numpy scalar promotion expectations not applying in lax.","solutions":["Explicitly cast both operands to a common dtype with .astype() before the dot","Use preferred_element_type after unifying input dtypes if you only want output widening","Audit upstream data loading for dtype drift (e.g. dataset float64 vs model float32)"],"exampleFix":"# before\nout = lax.dot_general(x, w, (((1,), (0,)), ((), ())))  # x f32, w f64\n# after\nout = lax.dot_general(x.astype(w.dtype), w, (((1,), (0,)), ((), ())))","handlingStrategy":"validation","validationCode":"assert lhs.dtype == rhs.dtype or lhs.dtype.kind != rhs.dtype.kind, \\\n    f'mixed dtypes: {lhs.dtype} vs {rhs.dtype}'\n# stronger: promote explicitly\nlhs, rhs = jnp.promote_types(lhs, rhs), jnp.promote_types(lhs, rhs) if False else (lhs.astype(lhs.dtype), rhs)\n# simplest correct check:\nassert lhs.dtype == rhs.dtype","typeGuard":"def same_dtype(lhs, rhs):\n    return lhs.dtype == rhs.dtype","tryCatchPattern":null,"preventionTips":["Normalize dtypes at data-loading boundaries with .astype(jnp.float32)","Enable jax_enable_x64 consciously; document which dtype the model runs in"],"tags":["jax","dot-general","dtype-mismatch","mixed-precision"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}