Create inspect_neural_networks.py

tools/inspect_neural_networks.py

@@ -0,0 +1,41 @@
+import torch
+import torch.nn as nn
+from smolagents import tool
+
+@tool
+def ml_layer_analyzer(layer_type: str, input_shape: list, **kwargs) -> str:
+    """A tool that dynamically instantiates a PyTorch layer, simulates a forward pass, 
+    and reports output shape and parameter count. Ideal for planning ML code.
+    
+    Args:
+        layer_type: The exact string name of the PyTorch layer (e.g., 'Conv2d', 'Linear', 'LSTM').
+        input_shape: A list of integers representing the input tensor dimensions (e.g., [1, 3, 224, 224]).
+        **kwargs: Arbitrary keyword arguments needed to configure the layer (e.g., in_features=10, out_features=20, out_channels=64, kernel_size=3).
+    """
+    try:
+        # Resolve layer class from torch.nn
+        if not hasattr(nn, layer_type):
+            return f"Error: '{layer_type}' is not a valid layer type in torch.nn"
+            
+        layer_cls = getattr(nn, layer_type)
+        layer_instance = layer_cls(**kwargs)
+        
+        # Create a mock tensor based on input shape
+        mock_input = torch.randn(*input_shape)
+        
+        # Simulate forward pass
+        with torch.no_grad():
+            output = layer_instance(mock_input)
+            
+        # Handle unpacking if the output is a tuple (like RNNs/LSTMs)
+        if isinstance(output, tuple):
+            out_shape = str([list(o.shape) if hasattr(o, 'shape') else type(o) for o in output])
+        else:
+            out_shape = list(output.shape)
+            
+        param_count = sum(p.numel() for p in layer_instance.parameters())
+        
+        return f"Success! Layer: {layer_type} | Output Shape: {out_shape} | Total Parameters: {param_count}"
+        
+    except Exception as e:
+        return f"Failed to analyze layer configuration. Error encountered: {str(e)}"