{"record":{"id":"e98839b74cfbe316","repo":"jax-ml/jax","slug":"preferred-element-type-must-not-be-narrower-than-e98839","errorCode":null,"errorMessage":"`preferred_element_type` must not be narrower than the original type, got preferred_element_type of {preferred_element_type} for result type of {result_dtype}.","messagePattern":"`preferred_element_type` must not be narrower than the original type, got preferred_element_type of (.+?) for result type of (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5898,"sourceCode":"    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)):\n    raise TypeError(\"`preferred_element_type` must not be narrower than the \"\n                    \"original type, got preferred_element_type of \"\n                    f\"{preferred_element_type} for result type of \"\n                    f\"{result_dtype}.\")\n  return preferred_element_type\n\ndef _dot_general_transpose_lhs(g, x, y, *, dimension_numbers, precision,\n                               preferred_element_type: DTypeLike | None,\n                               out_sharding, swap_ans=False):\n  (x_contract, y_contract), (x_batch, y_batch) = dimension_numbers\n  x_ndim = x.aval.ndim\n  x_kept = remaining(range(x_ndim), x_contract, x_batch)\n  y_kept = remaining(range(np.ndim(y)), y_contract, y_batch)\n  if swap_ans:\n    ans_batch, ans_y, _ = ranges_like(x_batch, y_kept, x_kept)\n  else:\n    ans_batch, _, ans_y = ranges_like(x_batch, x_kept, y_kept)\n  dims = ((ans_y, y_kept), (ans_batch, y_batch))\n  x_contract_sorted_by_y = list(np.take(x_contract, np.argsort(y_contract)))","sourceCodeStart":5880,"sourceCodeEnd":5916,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5880-L5916","documentation":"Raised by _maybe_upcast (replicating XLA's MaybeUpcast) when preferred_element_type has a smaller bit width than the natural result type and the result is not floating — e.g. asking for int8 output from int32 inputs. Non-float results cannot be safely narrowed.","triggerScenarios":"lax.dot_general(..., preferred_element_type=jnp.int8) with integer operands of wider dtype; also paths without a DotAlgorithm (check_bit_width True).","commonSituations":"Quantization-style code trying to force a narrow accumulator type; confusing preferred_element_type (output/accumulator type) with input casting.","solutions":["Choose preferred_element_type at least as wide as the input/result integer type","Cast inputs down first (x.astype(jnp.int8)) instead of narrowing the output","With DotAlgorithm, specify input/output types in the algorithm instead"],"exampleFix":"# before\nout = lax.dot_general(a, b, dn, preferred_element_type=jnp.int8)  # a,b int32\n# after\nout = lax.dot_general(a.astype(jnp.int8), b.astype(jnp.int8), dn,\n                      preferred_element_type=jnp.int32)","handlingStrategy":"validation","validationCode":"import numpy as np\nif preferred_element_type is not None:\n    w_in = np.dtype(preferred_element_type).itemsize * 8\n    w_out = np.dtype(result_dtype).itemsize * 8\n    assert result_dtype.kind == 'f' or w_in >= w_out","typeGuard":"def upcast_ok(result_dtype, pet):\n    import numpy as np\n    return pet is None or result_dtype == pet or (result_dtype.kind == 'f' or np.dtype(pet).itemsize >= np.dtype(result_dtype).itemsize)","tryCatchPattern":null,"preventionTips":["Remember preferred_element_type sets accumulator/output type, not a cast","Cast inputs explicitly for quantized paths"],"tags":["jax","dot-general","preferred-element-type","dtype-narrowing"],"backgroundTag":"dtype-narrowing-not-allowed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}