Residual typing issues in ULD

[lockwo]

#542 seems to still have some residual typing issues in dev. Running the checks on a fresh dev branch yields

pre-commit run --all-files
ruff-format..............................................................Passed
ruff.....................................................................Passed
pyright..................................................................Failed
- hook id: pyright
- exit code: 1

WARNING: there is a new pyright version available (v1.1.350 -> v1.1.392.post0).
Please install the new version or set PYRIGHT_PYTHON_FORCE_VERSION to `latest`

/diffrax/_solver/foster_langevin_srk.py
  /diffrax/_solver/foster_langevin_srk.py:66:64 - error: Expression of type "tuple[PyTree, PyTree, PyTree, PyTree, (UnderdampedLangevinX, PyTree) -> UnderdampedLangevinX]" cannot be assigned to return type "tuple[PyTree, PyTree, PyTree, PyTree, (UnderdampedLangevinX) -> UnderdampedLangevinX]"
    Type "(UnderdampedLangevinX, PyTree) -> UnderdampedLangevinX" cannot be assigned to type "(UnderdampedLangevinX) -> UnderdampedLangevinX" (reportReturnType)
/test/test_underdamped_langevin.py
 /test/test_underdamped_langevin.py:260:63 - error: Argument of type "(x: Unknown) -> Unknown" cannot be assigned to parameter "grad_f" of type "(UnderdampedLangevinX, PyTree) -> UnderdampedLangevinX" in function "__init__"
    Type "(x: Unknown) -> Unknown" cannot be assigned to type "(UnderdampedLangevinX, PyTree) -> UnderdampedLangevinX"
      Function accepts too many positional parameters; expected 1 but received 2 (reportArgumentType)
2 errors, 0 warnings, 0 informations

Also, I think the examples will fail since

drift_term = diffrax.UnderdampedLangevinDriftTerm(gamma, u, lambda x: 2 * x) is no longer correct (since the Args is required?)