Building jax from source without ".dev.*" in the final artifact version

[OsamaAlAhmadieh]

Hello,
probably a newbie question...
How could I specify the version that I want the final jaxlib wheel artifact to have when building jaxlib from source. I need to build it from source to remove AVX from jaxlib, but I always end up with .dev.* in the version, and it messes up when I try to add it to my project using poetry.. since in my project I have jax-0.6.2 as dependency.. and jaxlib needs to be exactly the same.

I couldn't find any specific flag in the build.py script, nor in the bazel-options.

would really appreciate any help.

Thank you!

[jakevdp]

Hello - the logic for this can be found here:

jax/jax/version.py

Lines 72 to 84 in 858efa0

	def _get_version_for_build() -> str:
	"""Determine the version at build time.
	The returned version string depends on which environment variables are set:
	- if WHEEL_VERSION_SUFFIX is set: version looks like "0.5.1.dev20230906+ge58560fdc"
	Here the WHEEL_VERSION_SUFFIX value is ".dev20230906+ge58560fdc".
	Please note that the WHEEL_VERSION_SUFFIX value is not the same as the
	JAX_CUSTOM_VERSION_SUFFIX value, and WHEEL_VERSION_SUFFIX is set by Bazel
	wheel build rule.
	- if JAX_RELEASE or JAXLIB_RELEASE are set: version looks like "0.4.16"
	- if JAX_NIGHTLY or JAXLIB_NIGHTLY are set: version looks like "0.4.16.dev20230906"
	- if none are set: version looks like "0.4.16.dev20230906+ge58560fdc
	"""

If the JAX_RELEASE or JAXLIB_RELEASE environment variable is set, the built version will be a release version without any dev or git tag info appended.

[OsamaAlAhmadieh]

Thank you !