jax-ml/jax · error · TypeError

reduce_window jvp does not support non-zero init_value_tange

Error message

reduce_window jvp does not support non-zero init_value_tangent.

What it means

The JVP rule for lax.reduce_window requires the tangent of the init_value to be zero (init_value is treated as a constant). If autodiff produces a non-zero tangent flowing into init_value, differentiation fails.

Source

Thrown at jax/_src/lax/windowed_reductions.py:444

def reduce_window_jvp(
    primals,
    tangents,
    window_dimensions,
    window_strides,
    padding,
    base_dilation,
    window_dilation,
    jaxpr,
    consts,
):

  reduction_jaxpr = jaxpr

  n = len(primals) // 2  # number of primal operands
  operand, init_value = util.split_list(primals, [n])
  operand_tangent, init_value_tangent = util.split_list(tangents, [n])
  if not all(isinstance(t, ad.Zero) for t in init_value_tangent):
    raise TypeError("reduce_window jvp does not support non-zero init_value_tangent.")

  init_value_tangent = map(ad_util.instantiate, init_value_tangent)
  c_reduction_jaxpr = reduction_jaxpr.with_consts(consts)
  jvp_reduction = ad.jvp_jaxpr(c_reduction_jaxpr, (True,) * len(tangents), [False] * len(init_value_tangent))[0]

  def wrapper(left, right):
    pl, tl = util.split_list(left, [n])
    pr, tr = util.split_list(right, [n])
    return jaxpr_as_fun(jvp_reduction)(*pl, *pr, *tl, *tr)

  jvp_primals_tangents = _reduce_window(
      operand=[*operand, *operand_tangent],
      init_value=[*init_value, *init_value_tangent],
      computation=wrapper,
      window_dimensions=window_dimensions,
      window_strides=window_strides,
      padding=padding,
      base_dilation=base_dilation,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Detach/stop-gradient the init_value: use jax.lax.stop_gradient(init_value) or a constant
2. Recompute init_value outside the traced differentiable path
3. Replace with a supported formulation, e.g. reduce_window_sum/max with fixed identity

Example fix

# before
out = lax.reduce_window(x, w, jaxpr, consts, dims, strides, padding)
# after
out = lax.reduce_window(x, jax.lax.stop_gradient(w), jaxpr, consts, dims, strides, padding)

init = jax.lax.stop_gradient(init_value)  # ensure zero tangent

When it happens

Trigger: Using jax.grad/jax.jvp on a function where the reduce_window init_value depends on a differentiated input (e.g. init_value = learned parameter * 0 + something non-zero tangent).

Common situations: Learned pooling identity parameters; init_value computed from trainable weights inside a pooling layer.

Related errors

• primal and tangent arguments to jax.jvp must be tuples or li
• primal and tangent arguments to jax.jvp must have the same t
• primal and tangent arguments to jax.jvp do not match; dtypes
• jvp called with different primal and tangent shapes;Got prim
• Pure callbacks do not support JVP. Please use `jax.custom_jv

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