quantax.nn.filter_grad

quantax.nn.filter_grad(fun: Callable, *, has_aux: bool = False, **gradkwargs) → Callable | Tuple[Callable, Any]

Creates a function that computes the gradient of fun similar to equinox.filter_grad.

The leaves in NoGradLayer are not considered as differentiable parameters.