[AHMETHAKANBEZIR1]

The custom norm JVP changes the primal value at zero: jax.jvp(norm, (zeros(3),), (ones(3),)) returns sqrt(3) rather than 0. This also affects the public Euclidean cost: value_and_grad(Euclidean()(x, fixed_x)) returns sqrt(3) when x equals fixed_x, although its ordinary value is 0.

The zero guard also tests the entire input instead of the requested norm axes. For rows [[0,0],[3,4]], the first row therefore has NaN JVP/reverse gradients when the second row is nonzero.

Keep the primal output from the original norm, and apply the clean-input/zero-tangent rule separately over each reduced slice. Preserve keepdims when constructing the broadcast mask, then squeeze only the reduced axes for the output mask.

Validation (Windows, Python3.12.14/JAX0.11.2 CPU):

• Sixteen new eager/JIT cases cover a zero vector, mixed zero/nonzero rows and columns, batched matrix Frobenius norms, and both keepdims modes. All16 fail on unmodified main36338c9 and pass after the fix, checking primal values, analytic JVPs, and reverse gradients.
• Entire math_utils_test.py:19 passed, including existing nonzero Hessian controls.
• tests/math plus tests/geometry/costs_test.py:73 passed,10 skipped.
• Public Euclidean value_and_grad smoke check: baseline sqrt(3)/zero gradient; fixed0/zero gradient.
• Configured pre-commit hooks and git diff --check pass.

Full package tests, GPU, and benchmarks were not run. Existing JAXopt deprecation warning remains. Prepared and tested autonomously by an AI coding agent; no independent human code review has been performed.