filter_value_and_grad: preserve return type of wrapped function

[ConnorBaker]

I ran into this while trying to update the RNN tutorial to use jaxtyping.

Using a bounded type variable conveys to the type checker that wrapping a function with filter_value_and_grad will not change the type of value returned by the original function.

As an example:

Suppose f returns a float. Let wrapped = filter_value_and_grad(f). Then the return type of wrapped is tuple[_Scalar, PyTree], not tuple[float, PyTree], where _Scalar is defined here:

equinox/equinox/_ad.py

Line 102 in 56bafcb

_Scalar = Union[float, complex, Float[ArrayLike, ""], Complex[ArrayLike, ""]]

The wrapped function can return a value of a different type so long as it is also a _Scalar (like complex).

Using the type variable communicates this invariant to the type checker.

I opted not to replace _Scalar with _ScalarTy throughout the entire file because it doesn't make sense to use a type variable which occurs exclusively arguments or the return value, and there were a few other use sites which do that with _Scalar.

[patrick-kidger]

Hmm, looks like checks are failing just for formatting. (See the contributing guide for how to have these checks run automatically locally, including automatically formatting.)

Other than that, this LGTM -- thank you for the fix.

[ConnorBaker]

That’s odd, I had installed and run them locally to make sure. I wonder if I have some global config set somewhere…

EDIT: Nope, just forgot that pre-commit looks at staged files. Fixed it.

Although, when I ran black . I did notice a number of files were re-formatted :x

[patrick-kidger]

Alright, LGTM! Merged. Thank you for contributing.