[envpool] Restore JAX XLA and keep Python 3.10 installs

Author: Trinkle23897

envpool/atari/atari_envpool_test.py

@@ -171,10 +171,7 @@ def test_xla_api(self) -> None:
       num_threads=2,
       thread_affinity_offset=0,
     )
-    try:
-      handle, recv, send, step = env.xla()
-    except RuntimeError as exc:
-      self.skipTest(str(exc))
+    handle, recv, send, step = env.xla()
     env.async_reset()
     handle, states = recv(handle)
     info = states[-1]
@@ -209,10 +206,7 @@ def test_xla_correctness(self) -> None:
       num_threads=2,
       thread_affinity_offset=0,
     )
-    try:
-      handle, recv, send, step = env1.xla()
-    except RuntimeError as exc:
-      self.skipTest(str(exc))
+    handle, recv, send, step = env1.xla()
     env1.async_reset()
     env2.async_reset()
 

envpool/python/xla_template.py

@@ -14,107 +14,69 @@
 """xla template on python side."""
 
 from collections import namedtuple
-from functools import partial
 from typing import Any, Callable, List, Tuple, Union, cast
 
 import numpy as np
-from jax import core, dtypes
+from jax import ShapeDtypeStruct, dtypes, ffi
 from jax import numpy as jnp
-from jax._src.lib import xla_client
-from jax.core import ShapedArray
-from jax.interpreters import xla
-
-_core = cast(Any, core)
-_xla = cast(Any, xla)
-_xla_client = cast(Any, xla_client)
-
-
-def _shape_with_layout(
-  specs: Tuple[Tuple[List[int], Any], ...]
-) -> Tuple[xla_client.Shape, ...]:
-  return tuple(
-    xla_client.Shape
-    .array_shape(dtype, shape, tuple(range(len(shape) -
-                                           1, -1, -1))) if len(shape) >
-    0 else xla_client.Shape.scalar_shape(dtype) for shape, dtype in specs
-  )
 
 
 def _normalize_specs(
   specs: Tuple[Tuple[Any, List[int]], ...]
-) -> Tuple[Tuple[List[int], Any], ...]:
+) -> Tuple[Tuple[Tuple[int, ...], Any], ...]:
   return tuple(
-    (shape, dtypes.canonicalize_dtype(dtype)) for dtype, shape in specs
+    (tuple(shape), dtypes.canonicalize_dtype(dtype)) for dtype, shape in specs
   )
 
 
+def _shape_dtype_struct(shape: Tuple[int, ...], dtype: Any) -> ShapeDtypeStruct:
+  return ShapeDtypeStruct(shape, dtype)
+
+
+def _layout(shape: Tuple[int, ...]) -> Tuple[int, ...]:
+  return tuple(range(len(shape)))
+
+
 def _make_xla_function(
   obj: Any,
   handle: bytes,
   name: str,
   specs: Tuple[Tuple[Any, ...], Tuple[Any, ...]],
   capsules: Tuple[Any, Any],
 ) -> Callable:
-  if not (
-    hasattr(_core, "Primitive") and
-    hasattr(_xla, "backend_specific_translations")
-  ):
-    raise RuntimeError(
-      "XLA is unavailable because this JAX version removed the legacy "
-      "primitive/translation APIs used by envpool."
-    )
   in_specs, out_specs = specs
   in_specs = _normalize_specs(in_specs)
   out_specs = _normalize_specs(out_specs)
   cpu_capsule, gpu_capsule = capsules
-  _xla_client.register_custom_call_target(
-    f"{type(obj).__name__}_{id(obj)}_{name}_cpu".encode(),
+  call_target_name = f"{type(obj).__name__}_{id(obj)}_{name}"
+  ffi.register_ffi_target(
+    call_target_name,
     cpu_capsule,
-    platform="cpu"
+    platform="cpu",
+    api_version=0,
   )
-  _xla_client.register_custom_call_target(
-    f"{type(obj).__name__}_{id(obj)}_{name}_gpu".encode(),
+  ffi.register_ffi_target(
+    call_target_name,
     gpu_capsule,
     platform="gpu",
+    api_version=0,
   )
-
-  def abstract(
-    *args: List[jnp.ndarray]
-  ) -> Union[ShapedArray, Tuple[ShapedArray, ...]]:
-    if len(out_specs) > 1:
-      return tuple(ShapedArray(*spec) for spec in out_specs)
-    else:
-      return ShapedArray(*out_specs[0])
-
-  def translation(c: Any, *args: Any, platform: str = "cpu") -> Any:
-    output_shape_with_layout = _shape_with_layout(out_specs)
-    if len(out_specs) == 1:
-      output_shape = output_shape_with_layout[0]
-    else:
-      output_shape = xla_client.Shape.tuple_shape(output_shape_with_layout)
-    return _xla_client.ops.CustomCallWithLayout(
-      c,
-      f"{type(obj).__name__}_{id(obj)}_{name}_{platform}".encode(),
-      operands=args,
-      operand_shapes_with_layout=_shape_with_layout(in_specs),
-      shape_with_layout=output_shape,
-      opaque=handle,
-      has_side_effect=True,
-    )
-
-  prim = _core.Primitive(f"{type(obj).__name__}_{id(obj)}_{name}")
-  prim.multiple_results = (len(out_specs) > 1)
-  prim.def_impl(partial(xla.apply_primitive, prim))
-  prim.def_abstract_eval(abstract)
-  _xla.backend_specific_translations["cpu"][prim] = partial(
-    translation, platform="cpu"
-  )
-  _xla.backend_specific_translations["gpu"][prim] = partial(
-    translation, platform="gpu"
+  result_specs = tuple(_shape_dtype_struct(*spec) for spec in out_specs)
+  xla_func = ffi.ffi_call(
+    call_target_name,
+    result_specs if len(result_specs) > 1 else result_specs[0],
+    has_side_effect=True,
+    input_layouts=tuple(_layout(shape) for shape, _ in in_specs),
+    output_layouts=(
+      tuple(_layout(shape) for shape, _ in out_specs)
+      if len(out_specs) > 1 else _layout(out_specs[0][0])
+    ),
+    custom_call_api_version=0,
+    legacy_backend_config=cast(Any, handle),
   )
 
   def call(*args: Any) -> Any:
-    return prim.bind(*args)
+    return xla_func(*args)
 
   return call
 

third_party/pip_requirements/requirements-dev-lock.txt

@@ -24,8 +24,10 @@ gymnasium==1.2.3
 h5py==3.16.0
 idna==3.11
 imageio==2.37.3
-jax==0.9.2
-jaxlib==0.9.2
+jax==0.6.2; python_version < "3.11"
+jax==0.9.2; python_version >= "3.11"
+jaxlib==0.6.2; python_version < "3.11"
+jaxlib==0.9.2; python_version >= "3.11"
 jinja2==3.1.6
 labmaze==1.0.6
 llvmlite==0.46.0

third_party/pip_requirements/requirements-dev.txt

@@ -5,7 +5,8 @@ dm-env>=1.6
 gym>=0.26.2
 gymnasium>=1.2.3
 optree>=0.19.0
-jax[cpu]>=0.9.2,<0.10
+jax[cpu]>=0.6.2,<0.7; python_version<'3.11'
+jax[cpu]>=0.9.2,<0.10; python_version>='3.11'
 absl-py
 packaging>=26.0
 tqdm

third_party/pip_requirements/requirements-release-lock.txt

@@ -7,8 +7,10 @@ farama-notifications==0.0.4
 gym==0.26.2
 gym-notices==0.1.0
 gymnasium==1.2.3
-jax==0.9.2
-jaxlib==0.9.2
+jax==0.6.2; python_version < "3.11"
+jax==0.9.2; python_version >= "3.11"
+jaxlib==0.6.2; python_version < "3.11"
+jaxlib==0.9.2; python_version >= "3.11"
 ml-dtypes==0.5.4
 numpy==2.4.3
 opt-einsum==3.4.0

third_party/pip_requirements/requirements-release.txt

@@ -5,5 +5,6 @@ dm-env>=1.6
 gym>=0.26.2
 gymnasium>=1.2.3
 optree>=0.19.0
-jax[cpu]>=0.9.2,<0.10
+jax[cpu]>=0.6.2,<0.7; python_version<'3.11'
+jax[cpu]>=0.9.2,<0.10; python_version>='3.11'
 packaging>=26.0