Use dynamically set version from __init__.py.

This allows programmatic querying of the Torax version, for eg. serialization. It is put in a separate version.py file to avoid circular dependencies.

PiperOrigin-RevId: 737640092

Author: sbodenstein

pyproject.toml

@@ -5,7 +5,6 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "torax"
 description = "Differentiable 1D tokamak plasma transport simulator in JAX."
-version = "0.3.1"
 readme = "README.md"
 requires-python = ">=3.10"
 license = {file = "LICENSE"}
@@ -42,6 +41,9 @@ dependencies = [
     "treelib>=1.3.2",
 ]
 
+# `version` is automatically set to use the `torax.__version__`
+dynamic = ["version"]
+
 [project.urls]
 homepage = "https://github.com/google-deepmind/torax"
 repository = "https://github.com/google-deepmind/torax"

torax/__init__.py

@@ -19,6 +19,7 @@
 
 # pylint: disable=g-importing-member
 
+from torax import version
 from torax.config.config_loader import import_module
 from torax.interpolated_param import InterpolatedVarSingleAxis
 from torax.interpolated_param import InterpolatedVarTimeRho
@@ -29,6 +30,8 @@
 
 # pylint: enable=g-importing-member
 
+__version__ = version.TORAX_VERSION
+__version_info__ = version.TORAX_VERSION_INFO
 
 __all__ = [
     'import_module',

torax/version.py

@@ -0,0 +1,29 @@
+# Copyright 2024 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""Torax version information."""
+
+from typing import Final
+
+# b/404741308: This is currently a closest tag that could be made more
+# fine-grained in the future, eg. based on an actual commit.
+TORAX_VERSION: Final[str] = "0.3.1"
+
+
+def _version_as_tuple(version_str: str) -> tuple[int, int, int]:
+  return tuple(int(i) for i in version_str.split(".") if i.isdigit())
+
+
+TORAX_VERSION_INFO: Final[tuple[int, int, int]] = _version_as_tuple(
+    TORAX_VERSION
+)