jax-ml/jax · error · TypeError

the first argument to pow must have an inexact dtype (float

Error message

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.

What it means

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.

Source

Thrown at jax/_src/lax/lax.py:4870

core.pp_eqn_rules[cbrt_p] = _unary_with_accuracy_pp_rule

square_p = standard_unop(_int | _float | _complex, 'square')

def _square_lower_hlo(ctx, x):
  if dtypes.issubdtype(ctx.avals_in[0].dtype, np.integer):
    return [hlo.multiply(x, x)]
  return [chlo.square(x)]

ad.defjvp2(square_p, lambda g, ans, x: mul(g, mul(_const(x, 2), x)))
mlir.register_lowering(square_p, _square_lower_hlo)

def _pow_dtype_rule(x, y):
  if (dtypes.issubdtype(x.dtype, np.inexact) and
      dtypes.issubdtype(y.dtype, np.integer)):
    return x.dtype
  if x.dtype == y.dtype:
    return x.dtype
  raise TypeError("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 "
                  f"{x.dtype} and {y.dtype} respectively.")
pow_p = naryop(_pow_dtype_rule, [_float | _complex, _int | _float | _complex],
               'pow', require_same_dtypes=False)

def _pow_jvp_lhs(g, ans, x, y):
  y_dtype = dtypes.dtype(y)
  result_dtype = dtypes.result_type(x, y)
  if result_dtype == bool:
    result_dtype = 'int32'
  x = convert_element_type(x, result_dtype)
  y = convert_element_type(y, result_dtype)
  if dtypes.issubdtype(y_dtype, np.integer):
    if x.shape != y.shape:
      shape = broadcast_shapes(x.shape, y.shape)
      sharding = broadcast_shardings(typeof(x), typeof(y))
      x = _maybe_broadcast(shape, x, sharding)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the base to an inexact dtype: x.astype(jnp.float32) ** y
2. Make dtypes match: promote both to the same float/complex dtype with jnp.promote_types or astype
3. For integer bases with integer exponents keep y integer; for roots use x.astype(float) ** 0.5 or jnp.sqrt

Example fix

// before
z = lax.pow(x_int, 0.5)
// after
z = lax.pow(x_int.astype(jnp.float32), 0.5)

x, y = jnp.promote_types(x, y), y.astype(jnp.promote_types(x, y))
if not jnp.issubdtype(x.dtype, jnp.inexact):
    x = x.astype(jnp.float32)
out = lax.pow(x, y)

def pow_dtypes_ok(x, y) -> bool:
    import numpy as np, jax.numpy as jnp
    return (np.issubdtype(x.dtype, np.inexact) and np.issubdtype(y.dtype, np.integer)) or x.dtype == y.dtype

When it happens

Trigger: 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.

Common situations: Porting numpy power semantics (numpy promotes more liberally); mixed precision (bfloat16 vs float32) exponentiation; int arrays raised to fractional powers intending roots.

Related errors

• full must be called with scalar fill_value, got fill_value.s
• offset must be an integer, got {offset!r}
• {} does not accept dtype {}. Accepted dtypes are subtypes of
• {} does not accept dtype {} at position {}. Accepted dtypes
• Integers cannot be raised to negative powers, got integer_po

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/e421cf05d5ac30e0. Report an issue: GitHub.