Isort formatting fixes.

PiperOrigin-RevId: 792640616

Author: Nush395

torax/_src/fvm/jax_root_finding.py

@@ -13,17 +13,15 @@
 # limitations under the License.
 
 """JAX root finding functions."""
-
-
 import dataclasses
 import functools
 from typing import Callable, Final
+
 import jax
 import jax.numpy as jnp
 import numpy as np
 from torax._src import jax_utils
 
-
 # Delta is a vector. If no entry of delta is above this magnitude, we terminate
 # the delta loop. This is to avoid getting stuck in an infinite loop in edge
 # cases with bad numerics.

torax/_src/fvm/tests/jax_root_finding_test.py

@@ -11,9 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
 import functools
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import chex

torax/tests/sim_experimental_compile_test.py

@@ -22,7 +22,6 @@
 from torax._src.output_tools import output
 from torax._src.test_utils import sim_test_case
 
-
 _ALL_PROFILES = (
     output.T_I,
     output.T_E,