{"record":{"id":"a9620f28f69bba00","repo":"jax-ml/jax","slug":"reduce-window-jvp-does-not-support-non-zero-init-v","errorCode":null,"errorMessage":"reduce_window jvp does not support non-zero init_value_tangent.","messagePattern":"reduce_window jvp does not support non-zero init_value_tangent\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/windowed_reductions.py","lineNumber":444,"sourceCode":"def reduce_window_jvp(\n    primals,\n    tangents,\n    window_dimensions,\n    window_strides,\n    padding,\n    base_dilation,\n    window_dilation,\n    jaxpr,\n    consts,\n):\n\n  reduction_jaxpr = jaxpr\n\n  n = len(primals) // 2  # number of primal operands\n  operand, init_value = util.split_list(primals, [n])\n  operand_tangent, init_value_tangent = util.split_list(tangents, [n])\n  if not all(isinstance(t, ad.Zero) for t in init_value_tangent):\n    raise TypeError(\"reduce_window jvp does not support non-zero init_value_tangent.\")\n\n  init_value_tangent = map(ad_util.instantiate, init_value_tangent)\n  c_reduction_jaxpr = reduction_jaxpr.with_consts(consts)\n  jvp_reduction = ad.jvp_jaxpr(c_reduction_jaxpr, (True,) * len(tangents), [False] * len(init_value_tangent))[0]\n\n  def wrapper(left, right):\n    pl, tl = util.split_list(left, [n])\n    pr, tr = util.split_list(right, [n])\n    return jaxpr_as_fun(jvp_reduction)(*pl, *pr, *tl, *tr)\n\n  jvp_primals_tangents = _reduce_window(\n      operand=[*operand, *operand_tangent],\n      init_value=[*init_value, *init_value_tangent],\n      computation=wrapper,\n      window_dimensions=window_dimensions,\n      window_strides=window_strides,\n      padding=padding,\n      base_dilation=base_dilation,","sourceCodeStart":426,"sourceCodeEnd":462,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/windowed_reductions.py#L426-L462","documentation":"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.","triggerScenarios":"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).","commonSituations":"Learned pooling identity parameters; init_value computed from trainable weights inside a pooling layer.","solutions":["Detach/stop-gradient the init_value: use jax.lax.stop_gradient(init_value) or a constant","Recompute init_value outside the traced differentiable path","Replace with a supported formulation, e.g. reduce_window_sum/max with fixed identity"],"exampleFix":"# before\nout = lax.reduce_window(x, w, jaxpr, consts, dims, strides, padding)\n# after\nout = lax.reduce_window(x, jax.lax.stop_gradient(w), jaxpr, consts, dims, strides, padding)","handlingStrategy":"validation","validationCode":"init = jax.lax.stop_gradient(init_value)  # ensure zero tangent","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Treat init_value as a constant in differentiable code","Run jax.jvp smoke tests on custom reductions"],"tags":["jax","autodiff","jvp","reduce-window"],"backgroundTag":"jax-autodiff-unsupported-differentiation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}