[DiadkoShmek]

What does this PR do?

Full graph-mode grad state filtering looks up only nonnegative argument positions. A valid negative argnums or DiffState therefore loses its module-state filter and can try to differentiate an integer BatchStat. The equivalent positive index correctly selects parameters and gives the expected gradient/state update.

Resolve the filter using both the positive position and its negative equivalent after arguments/defaults are resolved. The derivative, graph split/merge logic and public signatures are unchanged.

The new eager/JIT, grad/value_and_grad and negative-filter cases give 12 failures plus 12 positive controls passing on exact base; all 24 pass with the change. The selected TestGrad/bound-grad suite passes 63 tests including those regressions. Independent analytic module-gradient controls additionally cover auxiliary output, repeated state updates, kwargs/default gaps, dynamic arity, distinct negative filters and newly created graph state; invalid indices and positive/mode controls retain existing behavior.

Configured Ruff and whitespace checks pass. Validation uses real CPU NNX/JAX imports, not extracted graph methods; GPU/TPU, other JAX versions and the complete Flax suite were not tested. This does not change the separate intentional standalone-Variable or tree-mode behavior.

Checklist

• This PR fixes a small bug: one argument-filter lookup expression.
• This change includes high-coverage regression tests.
• Prior discussion: using the template's small-bug exception.
• Documentation: no public API/docstring change.

Review follow-up

Valid negative argument indices are now resolved to their argument slots per call before the duplicate guard and filter lookup. Aliases such as (1, -1) for two arguments are rejected, including distinct DiffState filters. Out-of-range indices retain the previous JAX error behavior; no modulo wrapping is used.

Validation: 48 alias regressions fail on the previous PR head and pass after this update; 105 TestGrad cases and 6 bound-grad controls pass. Separate checks cover 24 analytic filtered gradients, 14 unchanged error records and 24 variable-arity/default/keyword eager/JIT calls. Pinned Ruff 0.1.3 and whitespace checks pass. CPU JAX 0.12.0.1 only; full-project checks were not run locally.

[mangeshraut712]

Read the negative-argnums filter lookup and the new parameterized grad tests. The positive/negative fallback matches how JAX resolves argnums. One related hole in the duplicate guard below.

[mangeshraut712]

This fallback correctly finds a filter stored under -1 when iterating i in range(len(args)). The guard a few lines above still keys the dict with the raw index only (if index in index_filter), so argnums=(1, -1) for a 2-argument function is not treated as a repeated slot: both entries are stored, get(1, ...) wins for that argument, and jax_argnums still carries both 1 and -1.

Worth normalizing each index with index % len(args) (or index if index >= 0 else index + len(args)) before the duplicate check and dict insert, now that negative indices are supported for filtering.

[DiadkoShmek]

Fixed in e11687a. Valid negative indices are now resolved to argument slots after resolve_kwargs, before the duplicate check and filter lookup. Aliases such as (1, -1) and (-1, 1) are rejected even when they use distinct DiffState filters. Out-of-range indices keep the previous JAX error behavior.

Added regressions for both alias orders and both argument slots, eager/JIT and grad/value_and_grad. The 105 TestGrad cases plus 6 bound-grad controls pass; separate arity/default/keyword and unchanged-error checks pass too. Pinned Ruff 0.1.3 passes.