Update src/pybamm/solvers/idaklu_jax.py

Co-authored-by: Saransh Chopra <saransh0701@gmail.com>

Author: vidipsingh

src/pybamm/solvers/idaklu_jax.py

@@ -260,8 +260,8 @@ def f_isolated(*args, **kwargs):
     def jax_value(
         self,
         t: npt.NDArray[np.float64] | None = None,
-        inputs: Union[dict, None] = None,
-        output_variables: Union[list[str], None] = None,
+        inputs: dict | None = None,
+        output_variables: list[str] | None = None,
     ):
         """Helper function to compute the gradient of a jaxified expression