From 62f796bb3a1f1c1e30f15db18df9dd0eaeb7f612 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 08:10:44 +0000 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=B9=20Generalize=20=5Fcreate=5Fedges?= =?UTF-8?q?=5Ffrom=5Finputs=20in=20tracers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Moved the duplicated `_create_edges_from_inputs` logic from `TensorFlowTracer` and `PyTorchTracer` to the `BaseTracer`. Added `_is_framework_tensor` and `_tensor_to_metadata` abstract methods to ensure flexible integration of framework-specific checks. This removes duplication and simplifies maintainability across all tracer implementations. Co-authored-by: harshithluc073 <101515387+harshithluc073@users.noreply.github.com> --- neuroscope/core/tracer.py | 61 +++++++++++++++++++++++++++++++- neuroscope/tracers/jax.py | 21 +++++++++++ neuroscope/tracers/pytorch.py | 32 ++++------------- neuroscope/tracers/tensorflow.py | 33 ++++------------- 4 files changed, 93 insertions(+), 54 deletions(-) diff --git a/neuroscope/core/tracer.py b/neuroscope/core/tracer.py index 53b8dfc..ca83a7f 100644 --- a/neuroscope/core/tracer.py +++ b/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. diff --git a/neuroscope/tracers/jax.py b/neuroscope/tracers/jax.py index be6da2b..6781a8c 100644 --- a/neuroscope/tracers/jax.py +++ b/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 diff --git a/neuroscope/tracers/pytorch.py b/neuroscope/tracers/pytorch.py index e533c5c..8ebbb78 100644 --- a/neuroscope/tracers/pytorch.py +++ b/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 diff --git a/neuroscope/tracers/tensorflow.py b/neuroscope/tracers/tensorflow.py index ef67d7b..31793f1 100644 --- a/neuroscope/tracers/tensorflow.py +++ b/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] = {}