Move version.py to _src

pyproject.toml

@@ -44,7 +44,7 @@ dependencies = [
 dynamic = ["version"]
 
 [tool.setuptools.dynamic]
-version = {attr = "torax.version.TORAX_VERSION"}
+version = {attr = "torax._src.version.TORAX_VERSION"}
 
 [project.scripts]
 run_torax = "torax.run_simulation_main:run"

torax/__init__.py

@@ -20,7 +20,7 @@
 
 # pylint: disable=g-importing-member
 
-from torax import version
+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
 from torax._src.interpolated_param import InterpolatedVarSingleAxis

torax/_src/torax_pydantic/model_config.py

@@ -18,7 +18,7 @@
 import logging
 from typing import Any, Mapping
 import pydantic
-from torax import version
+from torax._src import version
 from torax._src.config import numerics as numerics_lib
 from torax._src.config import plasma_composition as plasma_composition_lib
 from torax._src.config import profile_conditions as profile_conditions_lib

torax/_src/torax_pydantic/tests/model_config_test.py

@@ -18,7 +18,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import chex
-from torax import version
+from torax._src import version
 from torax._src.config import config_loader
 from torax._src.torax_pydantic import model_config
 from torax._src.torax_pydantic import torax_pydantic