[TE/JAX] Update required JAX version for FFI custom calls with cudaGr…

…aph (#1285) Update jax version for ffi Signed-off-by: Phuong Nguyen <[email protected]>
NVIDIA · Oct 25, 2024 · 7cef756 · 7cef756
1 parent 7b284fe
commit 7cef756
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/transformer_engine/jax/cpp_extensions/misc.py b/transformer_engine/jax/cpp_extensions/misc.py
@@ -162,7 +162,7 @@ def is_ffi_enabled():
     """
     Helper function checking if XLA Custom Call with FFI is enabled
     """
-    is_supported = jax_version_meet_requirement("0.4.31")
+    is_supported = jax_version_meet_requirement("0.4.35")
     # New APIs with FFI are enabled by default
     is_enabled = int(os.getenv("NVTE_JAX_WITH_FFI", "1"))
     assert is_enabled in (0, 1), "Invalid NVTE_JAX_WITH_FFI value"