code: def patch_jax_version(version, release_version):
      """
      Patch jax.version._version & jax.version._release_version in order to
      test the version construction logic.
      """
      original_version = jax.version._version
      original_release_version = jax.version._release_version
  
      jax.version._version = version
      jax.version._release_version = release_version
      try:
          yield
      finally:
          jax.version._version = original_version
          jax.version._release_version = original_release_version
  
imports: import contextlib
  import datetime
  import unittest
  from unittest import mock
  import re
  from absl.testing import absltest
  import jax
  from jax._src.lib import check_jaxlib_version
  from jax._src import test_util
label_desc: Use image node_id 1 for calling a global function with example usage: patch_jax_version(version, release_version) without return types