{"record":{"id":"83b34aae69e4525d","repo":"jax-ml/jax","slug":"preferred-element-type-must-not-be-narrower-than","errorCode":null,"errorMessage":"`preferred_element_type` must not be narrower than the original type.","messagePattern":"`preferred_element_type` must not be narrower than the original type\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5679,"sourceCode":"    # different signedness between input and output.\n    pass\n  else:\n    allowed_types = (np.integer, np.floating, np.complexfloating)\n    if any(dtypes.issubdtype(input_dtype, t) and not\n           dtypes.issubdtype(preferred_element_type, t) for t in allowed_types):\n      raise TypeError(\"Input type is incompatible with \"\n                      \"`preferred_element_type`. The compatible combinations \"\n                      \"of (input_type, preferred_element_type) are \"\n                      \"(integral, integral), (integral, floating), \"\n                      \"(floating, floating), (complex, complex.\")\n    if (dtypes.issubdtype(input_dtype, np.signedinteger) and\n        not dtypes.issubdtype(preferred_element_type, np.signedinteger)):\n      raise TypeError(\"`preferred_element_type` must have the same signedness \"\n                      \"as the original type.\")\n  input_bitwidth = np.dtype(input_dtype).itemsize\n  preferred_bitwidth = np.dtype(preferred_element_type).itemsize\n  if preferred_bitwidth < input_bitwidth:\n    raise TypeError(\"`preferred_element_type` must not be narrower than the \"\n                    \"original type.\")\n\n\ndef _dot_general_shape_rule(lhs, rhs, *, dimension_numbers, precision,\n                            preferred_element_type: DTypeLike | None,\n                            out_sharding):\n  if out_sharding is not None and not isinstance(out_sharding, NamedSharding):\n    raise NotImplementedError\n  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = _from_maybe_ragged(dimension_numbers)\n  if not all(np.all(np.greater_equal(d, 0)) and np.all(np.less(d, lhs.ndim))\n             for d in (lhs_contracting, lhs_batch)):\n    msg = (\"dot_general requires lhs dimension numbers to be nonnegative and \"\n           \"less than the number of axes of the lhs value, got \"\n           f\"lhs_batch of {lhs_batch} and lhs_contracting of {lhs_contracting} \"\n           f\"for lhs of rank {lhs.ndim}\")\n    raise TypeError(msg)\n  if not all(np.all(np.greater_equal(d, 0)) and np.all(np.less(d, rhs.ndim))\n             for d in (rhs_contracting, rhs_batch)):","sourceCodeStart":5661,"sourceCodeEnd":5697,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5661-L5697","documentation":"dot_general's preferred_element_type (an accumulation/intermediate precision) must not be narrower in bits than the input dtype — you cannot request int8 accumulation for int16 inputs or float16 accumulation for float32 inputs, since that would lose precision silently. TypeError is raised when itemsize(preferred) < itemsize(input).","triggerScenarios":"lax.dot_general(a_float32, b_float32, ..., preferred_element_type=jnp.float16); int16 inputs with int8 preferred; also triggered via jnp.dot paths that forward precision settings on TPUs.","commonSituations":"Trying to speed up matmuls by requesting lower-precision accumulation; TF/TPU precision presets (bfloat16 semantics) applied to float32 inputs incorrectly; confusing preferred_element_type with output dtype.","solutions":["Use a preferred type with itemsize >= input itemsize (e.g. float32 or float64 for float32 inputs)","If lower precision is intended, downcast the inputs themselves first, then dot at that precision","For bf16 speedups, cast inputs to bfloat16 and use matching preferred type","Leave preferred_element_type=None to use input precision"],"exampleFix":"// before\nout = lax.dot_general(a_f32, b_f32, ..., preferred_element_type=jnp.float16)\n\n# after\na, b = a_f32.astype(jnp.bfloat16), b_f32.astype(jnp.bfloat16)\nout = lax.dot_general(a, b, ..., preferred_element_type=jnp.bfloat16)","handlingStrategy":"validation","validationCode":"if preferred is not None:\n    assert np.dtype(preferred).itemsize >= np.dtype(input_dtype).itemsize, 'too narrow'","typeGuard":"def wide_enough(input_dtype, preferred) -> bool:\n    return np.dtype(preferred).itemsize >= np.dtype(input_dtype).itemsize","tryCatchPattern":"try:\n    out = lax.dot_general(a, b, dn, preferred_element_type=pref)\nexcept TypeError:\n    out = lax.dot_general(a.astype(pref), b.astype(pref), dn, preferred_element_type=pref)","preventionTips":["Never request accumulation narrower than the inputs","Downcast inputs first if lower precision is the goal","Centralize precision policy in one helper that validates itemsizes"],"tags":["jax","dot-general","precision","dtype-validation"],"backgroundTag":"precision-narrowing-not-allowed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}