Design suggestion for better type check of has-aux

Currently, various optimizers accept a generic `Callable` function and parameters for specifying how it should be used:
```python
class LBFGS:
fun: Callable
value_and_grad: bool = False
has_aux: bool = False
```
So, the user has to carefully initialize making sure that these three variables are consistent.

Instead, it might be a better user experience to have:
```python
class LBFGS(Generic[T]):
fun: Callable[[T], RealNumeric] | ValueAuxAndGrad[T] | HasAux[T]
# No need for flags!
```
given the definitions:
```python
RealNumeric = jax.Array | npt.NDArray[np.floating[Any]] | float
T = TypeVar('T')

class _ValueAuxAndGradProtocol(Protocol, Generic[T]):
def __call__(self, x: T, /, *args: Any, **kwargs: Any
) -> tuple[tuple[RealNumeric, Any], T]:
...

class _HasAuxProtocol(Protocol, Generic[T]):
def __call__(self, x: T, /, *args: Any, **kwargs: Any
) -> tuple[RealNumeric, Any]:
...

@dataclass
class ValueAuxAndGrad(Generic[T]):
fun: _ValueAuxAndGradProtocol[T]
# Add appropriate methods based on how you use this

@dataclass
class HasAux(Generic[T]):
fun: _HasAuxProtocol[T]
# Add appropriate methods based on how you use this
```

This way, the function parmeters and return values are type-checked in all three cases. Also, there's no need for any flags. Instead of
```python
LBFGS(f, has_aux=True, value_and_grad=True)
```
you can do
```python
LBFGS(ValueAuxAndGrad(f)) # type-checked; fails if f doesn't return the appropriate tuple structure.
```
Even the simple use case would be type-checked:
```python
LBFGS(f) # type-checked; fails if f doesn't return an array or scalar float.
```

This would also allow transparent addition of `ValueAndGrad` one day, if it's desirable.

What do you think?