From b184495bad89b1d6f21ca0d486fa0f9766a868d5 Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Wed, 4 Mar 2026 13:18:50 +0000
Subject: [PATCH] Add tests for JAXTracer attach/detach and basic forward pass
 functionality.

Co-authored-by: harshithluc073 <101515387+harshithluc073@users.noreply.github.com>
---
 tests/test_jax_tracer.py | 78 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 78 insertions(+)
 create mode 100644 tests/test_jax_tracer.py

diff --git a/tests/test_jax_tracer.py b/tests/test_jax_tracer.py
new file mode 100644
index 0000000..b97eb0f
--- /dev/null
+++ b/tests/test_jax_tracer.py
@@ -0,0 +1,78 @@
+"""Tests for the JAX tracer."""
+
+import pytest
+
+
+class TestJAXTracer:
+    """Test suite for JAXTracer."""
+
+    @pytest.fixture
+    def simple_jax_model(self):
+        """Create a simple JAX function for testing."""
+        import jax.numpy as jnp
+
+        def my_function(x):
+            # Using simple operations that don't depend on specific shapes like dot
+            return x * 2.0 + jnp.ones_like(x)
+
+        return my_function
+
+    @pytest.fixture
+    def tracer(self):
+        """Create a JAXTracer instance."""
+        from neuroscope.tracers.jax import JAXTracer
+
+        # Disable error suppression to catch issues during testing
+        return JAXTracer(suppress_errors=False)
+
+    def test_attach_detach(self, tracer, simple_jax_model):
+        """Test attaching and detaching from a function."""
+        assert not tracer.is_attached
+
+        tracer.attach(simple_jax_model)
+        assert tracer.is_attached
+
+        tracer.detach()
+        assert not tracer.is_attached
+
+    def test_double_attach_raises(self, tracer, simple_jax_model):
+        """Test that attaching twice raises an error."""
+        tracer.attach(simple_jax_model)
+
+        with pytest.raises(RuntimeError, match="Already attached"):
+            tracer.attach(simple_jax_model)
+
+        tracer.detach()
+
+    def test_attach_invalid_type(self, tracer):
+        """Test that attaching to non-callable raises ValueError."""
+        with pytest.raises(ValueError, match="Expected callable"):
+            tracer.attach("not a function")
+
+    def test_forward_capture(self, tracer, simple_jax_model):
+        """Test that forward pass captures nodes and tensor metadata."""
+        import jax.numpy as jnp
+
+        traced_fn = tracer.attach(simple_jax_model)
+        tracer.reset_graph()
+
+        x = jnp.ones((4, 10))
+        _ = traced_fn(x)
+
+        graph = tracer.get_graph()
+
+        # Should have input nodes and operation nodes
+        assert len(graph.nodes) > 0
+
+        # Check that nodes have tensor metadata
+        has_tensors = False
+        for node in graph.nodes.values():
+            if node.output_tensors:
+                has_tensors = True
+                tensor_meta = node.output_tensors[0]
+                assert len(tensor_meta.shape) > 0
+                assert tensor_meta.dtype != ""
+
+        assert has_tensors, "No output tensors found in graph nodes"
+
+        tracer.detach()