#1104 tidy up

pybamm-team · Jul 13, 2020 · de4846e · de4846e
1 parent 1f691df
commit de4846e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/pybamm/solvers/jax_bdf_solver.py b/pybamm/solvers/jax_bdf_solver.py
@@ -13,7 +13,6 @@
 from jax.interpreters import partial_eval as pe
 from jax import linear_util as lu
 from jax.config import config
-from jax.lib import pytree
 
 config.update("jax_enable_x64", True)
 
@@ -688,8 +687,8 @@ def block_fun(i, j, Ai, Aj):
             )
 
     blocks = [
-        [ block_fun(i, j, Ai, Aj) for j, Aj in enumerate(lst)]
-            for i, Ai in enumerate(lst)
+        [block_fun(i, j, Ai, Aj) for j, Aj in enumerate(lst)]
+        for i, Ai in enumerate(lst)
     ]
 
     return jnp.block(blocks)

diff --git a/tests/unit/test_solvers/test_jax_bdf_solver.py b/tests/unit/test_solvers/test_jax_bdf_solver.py
@@ -92,7 +92,6 @@ def fun(y, t):
         np.testing.assert_allclose(y[:, 0], np.exp(0.1 * t_eval),
                                    rtol=1e-7, atol=1e-7)
 
-
     def test_solver_sensitivities(self):
         # Create model
         model = pybamm.BaseModel()