{"record":{"id":"e421cf05d5ac30e0","repo":"jax-ml/jax","slug":"the-first-argument-to-pow-must-have-an-inexact-dty","errorCode":null,"errorMessage":"the first argument to pow must have an inexact dtype (float or complex), and the second argument must have an inexact or integer dtype, and two inexact dtypes must match, but got {x.dtype} and {y.dtype} respectively.","messagePattern":"the first argument to pow must have an inexact dtype \\(float or complex\\), and the second argument must have an inexact or integer dtype, and two inexact dtypes must match, but got (.+?) and (.+?) respectively\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":4870,"sourceCode":"core.pp_eqn_rules[cbrt_p] = _unary_with_accuracy_pp_rule\n\nsquare_p = standard_unop(_int | _float | _complex, 'square')\n\ndef _square_lower_hlo(ctx, x):\n  if dtypes.issubdtype(ctx.avals_in[0].dtype, np.integer):\n    return [hlo.multiply(x, x)]\n  return [chlo.square(x)]\n\nad.defjvp2(square_p, lambda g, ans, x: mul(g, mul(_const(x, 2), x)))\nmlir.register_lowering(square_p, _square_lower_hlo)\n\ndef _pow_dtype_rule(x, y):\n  if (dtypes.issubdtype(x.dtype, np.inexact) and\n      dtypes.issubdtype(y.dtype, np.integer)):\n    return x.dtype\n  if x.dtype == y.dtype:\n    return x.dtype\n  raise TypeError(\"the first argument to pow must have an inexact dtype (float \"\n                  \"or complex), and the second argument must have an inexact or\"\n                  \" integer dtype, and two inexact dtypes must match, but got \"\n                  f\"{x.dtype} and {y.dtype} respectively.\")\npow_p = naryop(_pow_dtype_rule, [_float | _complex, _int | _float | _complex],\n               'pow', require_same_dtypes=False)\n\ndef _pow_jvp_lhs(g, ans, x, y):\n  y_dtype = dtypes.dtype(y)\n  result_dtype = dtypes.result_type(x, y)\n  if result_dtype == bool:\n    result_dtype = 'int32'\n  x = convert_element_type(x, result_dtype)\n  y = convert_element_type(y, result_dtype)\n  if dtypes.issubdtype(y_dtype, np.integer):\n    if x.shape != y.shape:\n      shape = broadcast_shapes(x.shape, y.shape)\n      sharding = broadcast_shardings(typeof(x), typeof(y))\n      x = _maybe_broadcast(shape, x, sharding)","sourceCodeStart":4852,"sourceCodeEnd":4888,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L4852-L4888","documentation":"lax.pow's dtype rule allows (inexact x, integer y) or equal dtypes. Any other mix — integer base with non-integer exponent, or mismatched float types like float32 ** float64 — raises this TypeError.","triggerScenarios":"lax.pow(int_array, 0.5); lax.pow(f32, f64); jnp powers of an int base with a negative or non-matching inexact exponent routed to pow; mixed-precision code passing bfloat16 ** float32.","commonSituations":"Porting numpy power semantics (numpy promotes more liberally); mixed precision (bfloat16 vs float32) exponentiation; int arrays raised to fractional powers intending roots.","solutions":["Cast the base to an inexact dtype: x.astype(jnp.float32) ** y","Make dtypes match: promote both to the same float/complex dtype with jnp.promote_types or astype","For integer bases with integer exponents keep y integer; for roots use x.astype(float) ** 0.5 or jnp.sqrt"],"exampleFix":"// before\nz = lax.pow(x_int, 0.5)\n// after\nz = lax.pow(x_int.astype(jnp.float32), 0.5)","handlingStrategy":"validation","validationCode":"x, y = jnp.promote_types(x, y), y.astype(jnp.promote_types(x, y))\nif not jnp.issubdtype(x.dtype, jnp.inexact):\n    x = x.astype(jnp.float32)\nout = lax.pow(x, y)","typeGuard":"def pow_dtypes_ok(x, y) -> bool:\n    import numpy as np, jax.numpy as jnp\n    return (np.issubdtype(x.dtype, np.inexact) and np.issubdtype(y.dtype, np.integer)) or x.dtype == y.dtype","tryCatchPattern":null,"preventionTips":["Promote operands to one float dtype before exponentiation","Never raise integer bases to fractional/negative powers; cast first"],"tags":["jax","lax","pow","dtype-mismatch","type-error"],"backgroundTag":"dtype-mismatch-in-binary-op","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}