Add configurable version.

This is to better support nightly builds, with custom version strings such as `0.1.0-dev20250305`.

e.g.
```
bazel run  --//jax_tpu_embedding/sparsecore:version_suffix=dev$(date '+%Y%m%d')  //build:build_pip_package --repo_env=HERMETIC_PYTHON_VERSION=3.12 -- $PWD
```

PiperOrigin-RevId: 733765952

jax_tpu_embedding/sparsecore/BUILD

@@ -11,16 +11,43 @@
 # 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.
+load("@bazel_skylib//rules:common_settings.bzl", "string_flag")
 load("//third_party/bazel/python:pytype.bzl", "pytype_strict_library")
+load(":configure_file.bzl", "configure_file")
 
 package(
     default_applicable_licenses = ["//:license"],
     default_visibility = ["//visibility:public"],
 )
 
+# Sets the version suffix for the package, e.g. dev20250321-805775f.
+string_flag(
+    name = "version_suffix",
+    build_setting_default = "",
+)
+
+# Sets the git commit for the package, e.g. 805775fcb5f9272e4c52dce751b00cf7f70364f2.
+string_flag(
+    name = "git_commit",
+    build_setting_default = "",
+)
+
+configure_file(
+    name = "version",
+    flag_substitutions = {
+        "VERSION_SUFFIX": ":version_suffix",
+        "GIT_COMMIT": ":git_commit",
+    },
+    output = "version.py",
+    template = "version.py.in",
+)
+
 pytype_strict_library(
     name = "sparsecore",
-    srcs = ["__init__.py"],
+    srcs = [
+        "__init__.py",
+        ":version",
+    ],
     deps = [
         "//jax_tpu_embedding/sparsecore/lib",  # buildcleaner: keep
         "//jax_tpu_embedding/sparsecore/utils",  # buildcleaner: keep

jax_tpu_embedding/sparsecore/__init__.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 """JAX SparseCore library."""
 
-# A new PyPI release will be pushed every time `__version__` is increased.
-# When changing this, also update the CHANGELOG.md.
-__version__ = '0.1.0'
-version = __version__
+from jax_tpu_embedding.sparsecore import version
+
+__version__ = version.__version__

jax_tpu_embedding/sparsecore/configure_file.bzl

@@ -0,0 +1,81 @@
+# Copyright 2024 The JAX SC Authors.
+#
+# 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.
+"""Generate a file using a template."""
+
+load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo")
+
+def _get_flag_substitutions(flag_substitutions):
+    """Extracts flag values."""
+    substitutions = {}
+    for key, label in flag_substitutions.items():
+        substitutions[key] = label[BuildSettingInfo].value
+    return substitutions
+
+def _create_substitution_map(string_substitutions):
+    """Replaces {key: value} with {${key}: value}"""
+    substitutions = {}
+    for key, value in string_substitutions.items():
+        key_var = "${" + key + "}"
+        substitutions[key_var] = value
+    return substitutions
+
+def configure_file(
+        name,
+        template,
+        output,
+        substitutions = {},
+        flag_substitutions = {}):
+    """Generates a file using a template.
+
+    For every entry in the substitutions maps, replaces `${variable}` with `value`.
+
+    Args:
+        name: The name of the rule.
+        template: The template file in which to perform the substitutions.
+        output: The output file.
+        substitutions: A map of string substitutions {variable: value}.
+        flag_substitutions: A map of variable to bazel string_flag substitutions. \
+            Replacement values are extracted from the flag.
+
+    Returns:
+        A rule that generates the output file.
+    """
+    _configure_file(
+        name = name,
+        template = template,
+        output = output,
+        substitutions = substitutions,
+        flag_substitutions = flag_substitutions,
+    )
+
+def _configure_file_impl(ctx):
+    substitutions = _create_substitution_map(ctx.attr.substitutions | _get_flag_substitutions(ctx.attr.flag_substitutions))
+    ctx.actions.expand_template(
+        template = ctx.file.template,
+        output = ctx.outputs.output,
+        substitutions = substitutions,
+    )
+
+_configure_file = rule(
+    implementation = _configure_file_impl,
+    attrs = {
+        "template": attr.label(
+            mandatory = True,
+            allow_single_file = True,
+        ),
+        "substitutions": attr.string_dict(),
+        "flag_substitutions": attr.string_keyed_label_dict(),
+        "output": attr.output(mandatory = True),
+    },
+)

jax_tpu_embedding/sparsecore/tests/BUILD

@@ -13,12 +13,23 @@
 # limitations under the License.
 load("//jax_tpu_embedding/sparsecore:tpu.bzl", "tpu_py_strict_test")
 load("//third_party/bazel/python:pypi.bzl", "pypi_requirement")
+load("//third_party/bazel/python:pytype.bzl", "pytype_strict_contrib_test")
 
 package(
     default_applicable_licenses = ["//:license"],
     default_visibility = ["//jax_tpu_embedding/sparsecore:__subpackages__"],
 )
 
+pytype_strict_contrib_test(
+    name = "version_test",
+    srcs = ["version_test.py"],
+    data = ["//jax_tpu_embedding/sparsecore:version"],
+    deps = [
+        "//jax_tpu_embedding/sparsecore",
+        pypi_requirement("absl/testing:absltest"),
+    ],
+)
+
 tpu_py_strict_test(
     name = "jax_sc_shakespeare_tests",
     srcs = ["jax_sc_shakespeare_tests.py"],

jax_tpu_embedding/sparsecore/tests/version_test.py

@@ -0,0 +1,30 @@
+# Copyright 2024 The JAX SC Authors.
+#
+# 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.
+"""Test version string generation."""
+
+from absl.testing import absltest
+from jax_tpu_embedding import sparsecore
+from jax_tpu_embedding.sparsecore import version
+
+
+class VersionTest(absltest.TestCase):
+
+  def test_version_string(self):
+    self.assertEqual(sparsecore.__version__, version.__version__)
+    self.assertTrue(version.__version__.startswith(version._base_version))
+    self.assertTrue(version.__version__.endswith(version._version_suffix))
+
+
+if __name__ == "__main__":
+  absltest.main()

jax_tpu_embedding/sparsecore/version.py.in

@@ -0,0 +1,18 @@
+"""JAX TPU Embedding versioning utilities
+
+For releases, the version is of the form:
+  xx.yy.zz
+
+For nightly builds, the date of the build is added:
+  xx.yy.zz-devYYYMMDD
+"""
+
+_base_version = "0.1.0"
+_version_suffix = "${VERSION_SUFFIX}"
+
+# Git commit corresponding to the build, if available.
+__git_commit__ = "${GIT_COMMIT}"
+
+# Library version.
+__version__ = _base_version + _version_suffix
+