def METHOD_NAME(self):
    # Create model
    model = pybamm.BaseModel()
    model.convert_to_format = "jax"
    domain = ["negative electrode", "separator", "positive electrode"]
    var = pybamm.Variable("var", domain=domain)
    model.rhs = {var: -pybamm.InputParameter("rate") * var}
    model.initial_conditions = {var: 1}
    # create discretisation
    mesh = get_mesh_for_testing(xpts=10)
    spatial_methods = {"macroscale": pybamm.FiniteVolume()}
    disc = pybamm.Discretisation(mesh, spatial_methods)
    disc.process_model(model)
    # Solve
    t_eval = np.linspace(0, 10, 4)
    y0 = model.concatenated_initial_conditions.evaluate().reshape(-1)
    rhs = pybamm.EvaluatorJax(model.concatenated_rhs)
    def fun(y, t, inputs):
        return rhs(t=t, y=y, inputs=inputs).reshape(-1)
    h = 0.0001
    rate = 0.1
    # create a dummy "model" where we calculate the sum of the time series
    @jax.jit
    def solve_bdf(rate):
        return jax.numpy.sum(
            pybamm.jax_bdf_integrate(
                fun, y0, t_eval, {"rate": rate}, rtol=1e-9, atol=1e-9
            )
        )
    # check answers with finite difference
    eval_plus = solve_bdf(rate + h)
    eval_neg = solve_bdf(rate - h)
    grad_num = (eval_plus - eval_neg) / (2 * h)
    grad_solve_bdf = jax.jit(jax.grad(solve_bdf))
    grad_bdf = grad_solve_bdf(rate)
    self.assertAlmostEqual(grad_bdf, grad_num, places=3)