Merge pull request #10148 from sharadmv:version-info

PiperOrigin-RevId: 439686412

jax/__init__.py

@@ -123,6 +123,7 @@
 )
 from jax.experimental.maps import soft_pmap as soft_pmap
 from jax.version import __version__ as __version__
+from jax.version import __version_info__ as __version_info__
 
 # These submodules are separate because they are in an import cycle with
 # jax and rely on the names imported above.

jax/version.py

@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+def _version_as_tuple(version_str):
+  return tuple(int(i) for i in version_str.split(".") if i.isdigit())
+
 __version__ = "0.3.5"
+__version_info__ = _version_as_tuple(__version__)
 
 _minimum_jaxlib_version = "0.3.0"
+_minimum_jaxlib_version_info = _version_as_tuple(_minimum_jaxlib_version)