Improved JAX detection

Author: pchanial

src/zeropybench/_jax.py

@@ -25,7 +25,7 @@ def from_code(cls, code: str, globals: dict[str, Any]) -> Self:
         return cls(tree, globals)
 
     def is_jax_context(self) -> bool:
-        """Returns true if a JAX variable or a jitted function is used."""
+        """Returns true if a JAX variable, module, or jitted function is used."""
         jax = sys.modules.get('jax')
         if jax is None:
             return False
@@ -39,8 +39,18 @@ def is_jax_context(self) -> bool:
                 return True
             if self._contains_jax_arrays(obj, jax):
                 return True
+            if self._is_jax_module(obj):
+                return True
         return False
 
+    @staticmethod
+    def _is_jax_module(obj: Any) -> bool:
+        """Check if an object is a JAX module (jax, jax.numpy, etc.)."""
+        module_name: str | None = getattr(obj, '__name__', None)
+        if module_name is None:
+            return False
+        return module_name == 'jax' or module_name.startswith('jax.')
+
     @staticmethod
     def _contains_jax_arrays(obj: Any, jax: Any) -> bool:
         """Check if an object is a pytree containing JAX arrays."""

tests/test_jax.py

@@ -121,6 +121,19 @@ class Config:
     assert parser.is_jax_context() is False
 
 
+def test_is_jax_context_jax_module() -> None:
+    """Test is_jax_context detects JAX modules (jnp.ones, jax.random, etc.)."""
+    # jax.numpy module
+    globals_ = {'jnp': jnp}
+    parser = CodeASTParser.from_code('jnp.ones((10, 10))', globals_)
+    assert parser.is_jax_context() is True
+
+    # jax module directly
+    globals_ = {'jax': jax}
+    parser = CodeASTParser.from_code('jax.numpy.ones((10, 10))', globals_)
+    assert parser.is_jax_context() is True
+
+
 def test_is_jax_context_no_jax(mocker: MockerFixture) -> None:
     """Test is_jax_context returns False when JAX is not imported."""
     mocker.patch.dict('sys.modules', {'jax': None})