pyink and isort

Author: mikesndrs

pyproject.toml

@@ -59,6 +59,7 @@ documentation = "https://github.com/google-deepmind/torax/blob/main/README.md"
 [project.optional-dependencies]
 # Installed through `pip install -e .[dev]`
 dev = [
+    "isort",
     "pytest",
     "pytest-xdist",
     "pytest-shard",
@@ -77,6 +78,9 @@ tutorial = [
     "notebook",
 ]
 
+[tool.isort]
+force_single_line = true
+
 [tool.pyink]
 # Formatting configuration to follow Google style-guide
 line-length = 80

torax/__init__.py

@@ -19,8 +19,6 @@
 
 import jax
 
-# pylint: disable=g-importing-member
-
 from torax._src import version
 from torax._src.config.config_loader import build_torax_config_from_file
 from torax._src.config.config_loader import import_module
@@ -35,6 +33,9 @@
 from torax._src.state import SolverNumericOutputs
 from torax._src.torax_pydantic.model_config import ToraxConfig
 
+# pylint: disable=g-importing-member
+
+
 # pylint: enable=g-importing-member
 
 __version__ = version.TORAX_VERSION
@@ -69,4 +70,5 @@ def set_jax_precision():
 def log_jax_backend():
   logging.info('JAX running on a default %s backend', jax.default_backend())
 
+
 set_jax_precision()

torax/_src/config/build_runtime_params.py

@@ -23,11 +23,12 @@
 DynamicRuntimeParamsSlice and a corresponding geometry with consistent Ip.
 """
 import chex
+import typing_extensions
+
 from torax._src.config import runtime_params_slice
 from torax._src.geometry import geometry
 from torax._src.geometry import geometry_provider as geometry_provider_lib
 from torax._src.torax_pydantic import model_config
-import typing_extensions
 
 
 def build_static_params_from_config(

torax/_src/config/config_loader.py

@@ -20,11 +20,13 @@
 import sys
 import types
 import typing
-from typing import Any, Literal, TypeAlias
+from typing import Any
+from typing import Literal
+from typing import TypeAlias
+
 from torax._src.plotting import plotruns_lib
 from torax._src.torax_pydantic import model_config
 
-
 ExampleConfig: TypeAlias = Literal[
     'basic_config', 'iterhybrid_predictor_corrector', 'iterhybrid_rampup'
 ]

torax/_src/config/numerics.py

@@ -15,9 +15,10 @@
 """Numerics parameters used throughout TORAX simulations."""
 import chex
 import pydantic
+from typing_extensions import Self
+
 from torax._src import array_typing
 from torax._src.torax_pydantic import torax_pydantic
-from typing_extensions import Self
 
 
 # pylint: disable=invalid-name

torax/_src/config/plasma_composition.py

@@ -17,6 +17,7 @@
 
 import chex
 import numpy as np
+
 from torax._src import array_typing
 from torax._src import constants
 from torax._src.config import runtime_validation_utils

torax/_src/config/profile_conditions.py

@@ -14,13 +14,17 @@
 
 """Profile condition parameters used throughout TORAX simulations."""
 import dataclasses
-from typing import Callable, Final
+from typing import Callable
+from typing import Final
+
 import chex
 import numpy as np
 import pydantic
+from typing_extensions import Self
+
 from torax._src import array_typing
 from torax._src.torax_pydantic import torax_pydantic
-from typing_extensions import Self
+
 # pylint: disable=invalid-name
 
 # Order of magnitude validations to catch common config errors.

torax/_src/config/runtime_params_slice.py

@@ -34,10 +34,12 @@
 restrictions as the dynamic arguments both in terms of types and how they are
 used.
 """
-from collections.abc import Mapping
 import dataclasses
+from collections.abc import Mapping
 
 import chex
+import typing_extensions
+
 from torax._src.config import numerics
 from torax._src.config import plasma_composition
 from torax._src.config import profile_conditions
@@ -50,7 +52,6 @@
 from torax._src.sources import runtime_params as sources_params
 from torax._src.torax_pydantic import torax_pydantic
 from torax._src.transport_model import runtime_params as transport_model_params
-import typing_extensions
 
 # Many of the variables follow scientific or mathematical notation, so disable
 # pylint complaints.
@@ -107,6 +108,7 @@ class StaticRuntimeParamsSlice:
   TODO(b/335596447): Add function to help users detect whether their
   change in config will trigger a recompile.
   """
+
   # Solver-specific static runtime params.
   solver: solver_params.StaticRuntimeParams
   # Mapping of source name to source-specific static runtime params.

torax/_src/config/runtime_validation_utils.py

@@ -14,17 +14,20 @@
 
 """Utilities for validating the config inputs."""
 
-from collections.abc import Mapping
 import functools
 import logging
-from typing import Annotated, Any, Final, TypeAlias
+from collections.abc import Mapping
+from typing import Annotated
+from typing import Any
+from typing import Final
+from typing import TypeAlias
 
 import numpy as np
 import pydantic
+
 from torax._src import constants
 from torax._src.torax_pydantic import torax_pydantic
 
-
 _TOLERANCE: Final[float] = 1e-6
 
 

torax/_src/config/tests/build_runtime_params_test.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 from absl.testing import absltest
 from absl.testing import parameterized
-import numpy as np
+
 from torax._src.config import build_runtime_params
 from torax._src.config import profile_conditions as profile_conditions_lib
 from torax._src.geometry import pydantic_model as geometry_pydantic_model