{"record":{"id":"2fe5dbd5147e4826","repo":"jax-ml/jax","slug":"ctx-avals-out-0-dtype","errorCode":null,"errorMessage":"{ctx.avals_out[0].dtype}","messagePattern":"\\{ctx\\.avals_out\\[0\\]\\.dtype\\}","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2864,"sourceCode":"  (aval_out,) = ctx.avals_out\n  out_type = ctx.aval_to_ir_type(aval_out)\n  assert isinstance(out_type, ir.ShapedType)\n  val_type = ir.ShapedType(out_type).element_type\n  if any(\n      isinstance(val_type, cls)\n      for cls in [\n          ir.BF16Type,\n          ir.F32Type,\n          ir.Float8E5M2Type,\n          ir.Float8E4M3FNType,\n          ir.Float8E4M3B11FNUZType,\n      ]\n  ):\n    val = ir.FloatAttr.get(val_type, 0.0)\n  elif isinstance(val_type, ir.IntegerType):\n    val = ir.IntegerAttr.get(val_type, 0)\n  else:\n    raise NotImplementedError(ctx.avals_out[0].dtype)\n  lhs_aval, rhs_aval = ctx.avals_in\n  # This is really a matrix-vector product. It only looks like matrix-matrix.\n  if (\n      lhs_dims == (1,)\n      and rhs_dims == (1,)\n      and ctx.avals_in[1].shape[0] == 1\n      and len(ctx.avals_in[0].shape) == 2\n      and len(ctx.avals_in[1].shape) == 2\n      and (\n          lhs_aval.dtype != jnp.float32\n          or rhs_aval.dtype != jnp.float32\n      )\n  ):\n    if ctx.avals_in[0].shape != ctx.avals_in[1].shape:\n      bcast_shape = jnp.broadcast_shapes(\n          ctx.avals_in[0].shape, ctx.avals_out[0].shape\n      )\n      bcast_shape = ir.VectorType.get(","sourceCodeStart":2846,"sourceCodeEnd":2882,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2846-L2882","documentation":"While constructing the zero-valued accumulator for a dot_general, Mosaic could not build an IR constant because the output dtype is neither a supported float nor an integer MLIR type (e.g. a complex or exotic dtype). The raw dtype is surfaced in the NotImplementedError message.","triggerScenarios":"dot_general (matmul) inside a Pallas Mosaic TPU kernel whose result/accumulator dtype is complex64 or another non-float/non-int type unsupported by the constant-building path.","commonSituations":"Complex-valued matmuls ported from CPU/GPU to Pallas TPU kernels.","solutions":["Compute the matmul in float and convert to complex afterwards, outside or after the dot","Split into real/imaginary float matmuls and recombine","Avoid complex dtypes in Pallas Mosaic kernels entirely"],"exampleFix":"// before\nz = jnp.matmul(a_c64, b_c64)\n// after\nre = jnp.matmul(a_c64.real.astype(jnp.float32), b_c64.real.astype(jnp.float32))\nim = jnp.matmul(a_c64.imag.astype(jnp.float32), b_c64.imag.astype(jnp.float32))","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nassert not jnp.issubdtype(out_dtype, jnp.complexfloating), 'complex dot unsupported in Mosaic'","typeGuard":"def is_dot_supported_dtype(dt) -> bool:\n    import jax.numpy as jnp\n    return jnp.issubdtype(dt, jnp.floating) or jnp.issubdtype(dt, jnp.integer)","tryCatchPattern":null,"preventionTips":["Restrict Pallas TPU kernel dtypes to float32/bf16/int32","Test kernels with each dtype you plan to feed them"],"tags":["jax","pallas","tpu","matmul","dtype","complex"],"backgroundTag":"unsupported-dtype-for-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}