🧹 Generalize _create_edges_from_inputs logic in tracers

neuroscope/core/tracer.py

@@ -13,7 +13,9 @@
 from typing import TYPE_CHECKING, Any, Callable
 
 if TYPE_CHECKING:
-    from neuroscope.core.graph import ExecutionGraph
+    from neuroscope.core.graph import ExecutionGraph, GraphEdge, TensorMetadata
+else:
+    from neuroscope.core.graph import GraphEdge, TensorMetadata
 
 
 class BaseTracer(ABC):
@@ -43,6 +45,9 @@ def __init__(self) -> None:
         self._last_broadcast_time = 0.0
         self._batch_lock = threading.Lock()
 
+        # Track tensor sources
+        self._tensor_sources: dict[int, str] = {}
+
     @property
     def is_attached(self) -> bool:
         """Whether the tracer is currently attached to a model."""
@@ -194,6 +199,60 @@ def on_forward_end(
         """
         pass
 
+    @abstractmethod
+    def _is_framework_tensor(self, tensor: Any) -> bool:
+        """
+        Check if a given object is a framework-specific tensor.
+
+        Args:
+            tensor: The object to check
+
+        Returns:
+            True if it's a framework tensor, False otherwise.
+        """
+        pass
+
+    @abstractmethod
+    def _tensor_to_metadata(self, tensor: Any) -> TensorMetadata:
+        """
+        Convert a framework-specific tensor to TensorMetadata.
+
+        Args:
+            tensor: Framework-specific tensor
+
+        Returns:
+            TensorMetadata for the tensor
+        """
+        pass
+
+    def _create_edges_from_inputs(self, inputs: Any, target_node_id: str) -> None:
+        """
+        Create edges from input tensor sources to this node.
+
+        Args:
+            inputs: The inputs to the module
+            target_node_id: The id of the current node
+        """
+        def process(tensor: Any, idx: int = 0) -> None:
+            if self._is_framework_tensor(tensor):
+                source_id = self._tensor_sources.get(id(tensor))
+                if source_id and source_id != target_node_id:
+                    edge = GraphEdge(
+                        source_id=source_id,
+                        target_id=target_node_id,
+                        target_input_idx=idx,
+                        tensor_info=self._tensor_to_metadata(tensor),
+                    )
+                    self._graph.add_edge(edge)
+            elif isinstance(tensor, (tuple, list)):
+                for i, t in enumerate(tensor):
+                    process(t, i)
+            elif isinstance(tensor, dict):
+                for i, t in enumerate(tensor.values()):
+                    process(t, i)
+
+        process(inputs)
+
     def _classify_module_type(self, module_class_name: str) -> str:
         """
         Classify a module into a category for styling purposes.

neuroscope/tracers/jax.py

@@ -129,6 +129,27 @@ def on_forward_end(
         """Not used for JAX tracing."""
         pass
 
+    def _is_framework_tensor(self, tensor: Any) -> bool:
+        """Check if an object is a JAX array."""
+        try:
+            import jax.numpy as jnp
+            import jax
+            return isinstance(tensor, (jnp.ndarray, jax.Array))
+        except ImportError:
+            return hasattr(tensor, "shape") and hasattr(tensor, "dtype")
+
+    def _tensor_to_metadata(self, tensor: Any) -> TensorMetadata:
+        """Convert JAX tensor to metadata."""
+        # JAX tracer primarily uses _aval_to_metadata, but we must implement the abstract method
+        if hasattr(tensor, "shape"):
+            return TensorMetadata(
+                shape=tuple(int(d) if hasattr(d, '__int__') else -1 for d in tensor.shape),
+                dtype=str(tensor.dtype) if hasattr(tensor, "dtype") else "unknown",
+                device="unknown",
+                requires_grad=False,
+            )
+        return TensorMetadata(shape=(), dtype="unknown")
+
     def _create_wrapper(self, fn: Callable) -> Callable:
         """Create a wrapped function that traces execution."""
         import jax

neuroscope/tracers/pytorch.py

@@ -86,7 +86,6 @@ def __init__(
         self._name_to_depth: dict[str, int] = {}
         self._name_to_parent: dict[str, str | None] = {}
         self._active_modules: list[str] = []  # Stack for tracking nested calls
-        self._tensor_sources: dict[int, str] = {}  # tensor id -> source node id
         self._suppress_errors = suppress_errors
         self._error_count = 0
 
@@ -193,7 +192,7 @@ def reset_graph(self) -> None:
         self._graph.clear()
         self._execution_order = 0
         self._active_modules = []
-        self._tensor_sources = {}
+        self._tensor_sources.clear()
         self._module_start_times = {}
         self._module_start_memory = {}
         self._name_to_node_id = {}
@@ -360,6 +359,11 @@ def _tensor_to_metadata(self, tensor: Any) -> TensorMetadata:
             is_contiguous=tensor.is_contiguous(),
         )
 
+    def _is_framework_tensor(self, tensor: Any) -> bool:
+        """Check if an object is a PyTorch tensor."""
+        import torch
+        return isinstance(tensor, torch.Tensor)
+
     def _check_tensor_issues(self, tensors: Any) -> bool:
         """Check for NaN or Inf values in tensors."""
         import torch
@@ -427,30 +431,6 @@ def _update_tensor_sources(self, outputs: Any, node_id: str) -> None:
             for t in outputs.values():
                 self._update_tensor_sources(t, node_id)
 
-    def _create_edges_from_inputs(self, inputs: Any, target_node_id: str) -> None:
-        """Create edges from input tensor sources to this node."""
-        import torch
-
-        def process(tensor: Any, idx: int = 0) -> None:
-            if isinstance(tensor, torch.Tensor):
-                source_id = self._tensor_sources.get(id(tensor))
-                if source_id and source_id != target_node_id:
-                    edge = GraphEdge(
-                        source_id=source_id,
-                        target_id=target_node_id,
-                        target_input_idx=idx,
-                        tensor_info=self._tensor_to_metadata(tensor),
-                    )
-                    self._graph.add_edge(edge)
-            elif isinstance(tensor, (tuple, list)):
-                for i, t in enumerate(tensor):
-                    process(t, i)
-            elif isinstance(tensor, dict):
-                for i, t in enumerate(tensor.values()):
-                    process(t, i)
-
-        process(inputs)
-
     def _compute_tensor_stats(self, tensors: Any) -> TensorStats | None:
         """Compute statistical summary of output tensors."""
         import torch

neuroscope/tracers/tensorflow.py

@@ -68,7 +68,6 @@ def __init__(
         self._layer_depths: dict[int, int] = {}
         self._suppress_errors = suppress_errors
         self._error_count = 0
-        self._tensor_sources: dict[int, str] = {}
 
         # v0.2.0 profiling
         self._enable_profiling = enable_profiling
@@ -101,7 +100,6 @@ def attach(self, model: Any) -> None:
         self._model = model
         self._original_calls = {}
         self._layer_depths = {}
-        self._tensor_sources = {}
 
         # Calculate layer depths
         self._calculate_layer_depths(model)
@@ -141,7 +139,7 @@ def reset_graph(self) -> None:
         """Clear the execution graph."""
         self._graph.clear()
         self._execution_order = 0
-        self._tensor_sources = {}
+        self._tensor_sources.clear()
         self._error_count = 0
 
     def on_forward_start(self, module: Any, inputs: Any, name: str) -> None:
@@ -315,6 +313,11 @@ def _tensor_to_metadata(self, tensor: Any) -> TensorMetadata:
             memory_bytes=0,  # TF doesn't expose this easily
         )
 
+    def _is_framework_tensor(self, tensor: Any) -> bool:
+        """Check if an object is a TensorFlow tensor."""
+        import tensorflow as tf
+        return isinstance(tensor, (tf.Tensor, tf.Variable))
+
     def _check_tensor_issues(self, tensors: Any) -> bool:
         """Check for NaN or Inf values in tensors."""
         import tensorflow as tf
@@ -344,30 +347,6 @@ def _update_tensor_sources(self, outputs: Any, node_id: str) -> None:
             for t in outputs.values():
                 self._update_tensor_sources(t, node_id)
 
-    def _create_edges_from_inputs(self, inputs: Any, target_node_id: str) -> None:
-        """Create edges from input tensor sources to this node."""
-        import tensorflow as tf
-
-        def process(tensor: Any, idx: int = 0) -> None:
-            if isinstance(tensor, (tf.Tensor, tf.Variable)):
-                source_id = self._tensor_sources.get(id(tensor))
-                if source_id and source_id != target_node_id:
-                    edge = GraphEdge(
-                        source_id=source_id,
-                        target_id=target_node_id,
-                        target_input_idx=idx,
-                        tensor_info=self._tensor_to_metadata(tensor),
-                    )
-                    self._graph.add_edge(edge)
-            elif isinstance(tensor, (tuple, list)):
-                for i, t in enumerate(tensor):
-                    process(t, i)
-            elif isinstance(tensor, dict):
-                for i, t in enumerate(tensor.values()):
-                    process(t, i)
-
-        process(inputs)
-
     def _get_layer_info(self, layer: Any) -> dict[str, Any]:
         """Extract layer configuration info."""
         info: dict[str, Any] = {}