Bump minimum jaxlib version to 0.3.2 and remove transfer guard compat…

…ibility code

build/build_wheel.py

@@ -99,7 +99,6 @@ def copy_file(src_file, dst_dir, dst_filename=None):
     "transfer_guard_lib.pyi",
 ]
 _OPTIONAL_XLA_EXTENSION_STUBS = [
-    "transfer_guard_lib.pyi",  # Will be required on xla_extension_version >= 58.
 ]
 
 

jax/_src/config.py

@@ -26,8 +26,7 @@
 
 from jax._src import lib
 from jax._src.lib import jax_jit
-if lib.xla_extension_version >= 58:
-  from jax._src.lib import transfer_guard_lib
+from jax._src.lib import transfer_guard_lib
 
 def bool_env(varname: str, default: bool) -> bool:
   """Read an environment variable and interpret it as a boolean.
@@ -714,139 +713,116 @@ def _update_disable_jit_thread_local(val):
     default=True,
     help=('Enables using optimization-barrier op for lowering remat.'))
 
-if lib.xla_extension_version < 58:
-  @contextlib.contextmanager
-  def explicit_device_put_scope() -> Iterator[None]:
-    """Indicates that the current context is an explicit device_put*() call."""
+@contextlib.contextmanager
+def explicit_device_put_scope() -> Iterator[None]:
+  """Indicates that the current context is an explicit device_put*() call."""
+  state = transfer_guard_lib.thread_local_state()
+  prev = state.explicit_device_put
+  state.explicit_device_put = True
+  try:
     yield
-
-  @contextlib.contextmanager
-  def explicit_device_get_scope() -> Iterator[None]:
-    """Indicates that the current context is an explicit device_get() call."""
+  finally:
+    state.explicit_device_put = prev
+
+@contextlib.contextmanager
+def explicit_device_get_scope() -> Iterator[None]:
+  """Indicates that the current context is an explicit device_get() call."""
+  state = transfer_guard_lib.thread_local_state()
+  prev = state.explicit_device_get
+  state.explicit_device_get = True
+  try:
     yield
-
-  @contextlib.contextmanager
-  def _transfer_guard(new_val: str) -> Iterator[None]:
-    raise NotImplementedError("jaxlib version is too low for transfer guards")
-
-  transfer_guard_host_to_device = _transfer_guard
-  transfer_guard_device_to_device = _transfer_guard
-  transfer_guard_device_to_host = _transfer_guard
-  transfer_guard = _transfer_guard
-
-else:
-  @contextlib.contextmanager
-  def explicit_device_put_scope() -> Iterator[None]:
-    """Indicates that the current context is an explicit device_put*() call."""
-    state = transfer_guard_lib.thread_local_state()
-    prev = state.explicit_device_put
-    state.explicit_device_put = True
-    try:
-      yield
-    finally:
-      state.explicit_device_put = prev
-
-  @contextlib.contextmanager
-  def explicit_device_get_scope() -> Iterator[None]:
-    """Indicates that the current context is an explicit device_get() call."""
-    state = transfer_guard_lib.thread_local_state()
-    prev = state.explicit_device_get
-    state.explicit_device_get = True
-    try:
-      yield
-    finally:
-      state.explicit_device_get = prev
-
-  def _update_transfer_guard(state, key, val):
-    """Applies the transfer guard level within transfer_guard_lib."""
-    if val is None:
-      setattr(state, key, None)
-    elif val == 'allow':
-      setattr(state, key, transfer_guard_lib.TransferGuardLevel.ALLOW)
-    elif val == 'log':
-      setattr(state, key, transfer_guard_lib.TransferGuardLevel.LOG)
-    elif val == 'disallow':
-      setattr(state, key, transfer_guard_lib.TransferGuardLevel.DISALLOW)
-    elif val == 'log_explicit':
-      setattr(state, key, transfer_guard_lib.TransferGuardLevel.LOG_EXPLICIT)
-    elif val == 'disallow_explicit':
-      setattr(state, key,
-              transfer_guard_lib.TransferGuardLevel.DISALLOW_EXPLICIT)
-    else:
-      assert False, f"Invalid transfer guard level {val}"
-
-  transfer_guard_host_to_device = config.define_enum_state(
-      name='jax_transfer_guard_host_to_device',
-      enum_values=[
-          'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit'
-      ],
-      # The default is applied by transfer_guard_lib. Use None here to avoid
-      # accidentally overriding --jax_transfer_guard.
-      default=None,
-      help=('Select the transfer guard level for host-to-device transfers. '
-            'Default is "allow".'),
-      update_global_hook=lambda val: _update_transfer_guard(
-          transfer_guard_lib.global_state(), 'host_to_device', val),
-      update_thread_local_hook=lambda val: _update_transfer_guard(
-          transfer_guard_lib.thread_local_state(), 'host_to_device', val))
-
-  transfer_guard_device_to_device = config.define_enum_state(
-      name='jax_transfer_guard_device_to_device',
-      enum_values=[
-          'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit'
-      ],
-      # The default is applied by transfer_guard_lib. Use None here to avoid
-      # accidentally overriding --jax_transfer_guard.
-      default=None,
-      help=('Select the transfer guard level for device-to-device transfers. '
-            'Default is "allow".'),
-      update_global_hook=lambda val: _update_transfer_guard(
-          transfer_guard_lib.global_state(), 'device_to_device', val),
-      update_thread_local_hook=lambda val: _update_transfer_guard(
-          transfer_guard_lib.thread_local_state(), 'device_to_device', val))
-
-  transfer_guard_device_to_host = config.define_enum_state(
-      name='jax_transfer_guard_device_to_host',
-      enum_values=[
-          'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit'
-      ],
-      # The default is applied by transfer_guard_lib. Use None here to avoid
-      # accidentally overriding --jax_transfer_guard.
-      default=None,
-      help=('Select the transfer guard level for device-to-host transfers. '
-            'Default is "allow".'),
-      update_global_hook=lambda val: _update_transfer_guard(
-          transfer_guard_lib.global_state(), 'device_to_host', val),
-      update_thread_local_hook=lambda val: _update_transfer_guard(
-          transfer_guard_lib.thread_local_state(), 'device_to_host', val))
-
-  def _update_all_transfer_guard_global(val):
-    for name in ('jax_transfer_guard_host_to_device',
-                 'jax_transfer_guard_device_to_device',
-                 'jax_transfer_guard_device_to_host'):
-      config.update(name, val)
-
-  _transfer_guard = config.define_enum_state(
-      name='jax_transfer_guard',
-      enum_values=[
-          'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit'
-      ],
-      # The default is applied by transfer_guard_lib. Use None here to avoid
-      # accidentally overriding --jax_transfer_guard_*.
-      default=None,
-      help=(
-          'Select the transfer guard level for all transfers. This option is '
+  finally:
+    state.explicit_device_get = prev
+
+def _update_transfer_guard(state, key, val):
+  """Applies the transfer guard level within transfer_guard_lib."""
+  if val is None:
+    setattr(state, key, None)
+  elif val == 'allow':
+    setattr(state, key, transfer_guard_lib.TransferGuardLevel.ALLOW)
+  elif val == 'log':
+    setattr(state, key, transfer_guard_lib.TransferGuardLevel.LOG)
+  elif val == 'disallow':
+    setattr(state, key, transfer_guard_lib.TransferGuardLevel.DISALLOW)
+  elif val == 'log_explicit':
+    setattr(state, key, transfer_guard_lib.TransferGuardLevel.LOG_EXPLICIT)
+  elif val == 'disallow_explicit':
+    setattr(state, key, transfer_guard_lib.TransferGuardLevel.DISALLOW_EXPLICIT)
+  else:
+    assert False, f'Invalid transfer guard level {val}'
+
+transfer_guard_host_to_device = config.define_enum_state(
+    name='jax_transfer_guard_host_to_device',
+    enum_values=[
+        'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit'
+    ],
+    # The default is applied by transfer_guard_lib. Use None here to avoid
+    # accidentally overriding --jax_transfer_guard.
+    default=None,
+    help=('Select the transfer guard level for host-to-device transfers. '
+          'Default is "allow".'),
+    update_global_hook=lambda val: _update_transfer_guard(
+        transfer_guard_lib.global_state(), 'host_to_device', val),
+    update_thread_local_hook=lambda val: _update_transfer_guard(
+        transfer_guard_lib.thread_local_state(), 'host_to_device', val))
+
+transfer_guard_device_to_device = config.define_enum_state(
+    name='jax_transfer_guard_device_to_device',
+    enum_values=[
+        'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit'
+    ],
+    # The default is applied by transfer_guard_lib. Use None here to avoid
+    # accidentally overriding --jax_transfer_guard.
+    default=None,
+    help=('Select the transfer guard level for device-to-device transfers. '
+          'Default is "allow".'),
+    update_global_hook=lambda val: _update_transfer_guard(
+        transfer_guard_lib.global_state(), 'device_to_device', val),
+    update_thread_local_hook=lambda val: _update_transfer_guard(
+        transfer_guard_lib.thread_local_state(), 'device_to_device', val))
+
+transfer_guard_device_to_host = config.define_enum_state(
+    name='jax_transfer_guard_device_to_host',
+    enum_values=[
+        'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit'
+    ],
+    # The default is applied by transfer_guard_lib. Use None here to avoid
+    # accidentally overriding --jax_transfer_guard.
+    default=None,
+    help=('Select the transfer guard level for device-to-host transfers. '
+          'Default is "allow".'),
+    update_global_hook=lambda val: _update_transfer_guard(
+        transfer_guard_lib.global_state(), 'device_to_host', val),
+    update_thread_local_hook=lambda val: _update_transfer_guard(
+        transfer_guard_lib.thread_local_state(), 'device_to_host', val))
+
+def _update_all_transfer_guard_global(val):
+  for name in ('jax_transfer_guard_host_to_device',
+               'jax_transfer_guard_device_to_device',
+               'jax_transfer_guard_device_to_host'):
+    config.update(name, val)
+
+_transfer_guard = config.define_enum_state(
+    name='jax_transfer_guard',
+    enum_values=[
+        'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit'
+    ],
+    # The default is applied by transfer_guard_lib. Use None here to avoid
+    # accidentally overriding --jax_transfer_guard_*.
+    default=None,
+    help=('Select the transfer guard level for all transfers. This option is '
           'set-only; the transfer guard level for a specific direction should '
           'be read using the per-transfer direction option. '
           'Default is "allow".'),
-      update_global_hook=_update_all_transfer_guard_global)
-
-  @contextlib.contextmanager
-  def transfer_guard(new_val: str) -> Iterator[None]:
-    """Set up thread-local state and return a contextmanager for managing it."""
-    with contextlib.ExitStack() as stack:
-      stack.enter_context(transfer_guard_host_to_device(new_val))
-      stack.enter_context(transfer_guard_device_to_device(new_val))
-      stack.enter_context(transfer_guard_device_to_host(new_val))
-      stack.enter_context(_transfer_guard(new_val))
-      yield
+    update_global_hook=_update_all_transfer_guard_global)
+
+@contextlib.contextmanager
+def transfer_guard(new_val: str) -> Iterator[None]:
+  """Set up thread-local state and return a contextmanager for managing it."""
+  with contextlib.ExitStack() as stack:
+    stack.enter_context(transfer_guard_host_to_device(new_val))
+    stack.enter_context(transfer_guard_device_to_device(new_val))
+    stack.enter_context(transfer_guard_device_to_host(new_val))
+    stack.enter_context(_transfer_guard(new_val))
+    yield

jax/_src/lib/__init__.py

@@ -173,5 +173,4 @@ def _parse_version(v: str) -> Tuple[int, ...]:
 if not os.path.isdir(cuda_path):
   cuda_path = None
 
-if xla_extension_version >= 58:
-  transfer_guard_lib = xla_client._xla.transfer_guard_lib
+transfer_guard_lib = xla_client._xla.transfer_guard_lib

jax/version.py

@@ -18,5 +18,5 @@ def _version_as_tuple(version_str):
 __version__ = "0.3.6"
 __version_info__ = _version_as_tuple(__version__)
 
-_minimum_jaxlib_version = "0.3.0"
+_minimum_jaxlib_version = "0.3.2"
 _minimum_jaxlib_version_info = _version_as_tuple(_minimum_jaxlib_version)