Following this instruction :
1. Strictly adhere to PEP-8 standards, ensuring code is clean, well-structured, and follows Python's style guide for readability.
2. Use appropriate modules like `typing` for type hints, ensuring code clarity, maintainability, and leveraging Python's standard libraries effectively.
3. Focus on time, space, robustness, scalability; write optimized, efficient code that can handle large datasets and scale seamlessly.
4. Implement comprehensive error handling, catching exceptions gracefully, providing meaningful error messages, and ensuring code robustness under unexpected conditions.
5. Ensure the developer agent writes advanced, maintainable Python code, balancing performance and readability while adhering to best practices.


TASK : use ploty module and torch given model , most generalised 

Develop a PyTorch model extending torch.nn.Module that accepts a recursive key input and returns a tensor of varying dimensions (1D, 2D, 3D, or 4D). Based on the dimensionality of the returned tensor, implement a corresponding scatter plot visualization. For 1D tensors, generate a simple scatter plot; for 2D tensors, create a standard 2D scatter plot. For 3D tensors, produce an advanced 3D scatter plot with interactive features. For 4D tensors, also create an advanced 3D scatter plot, incorporating a single legend that dynamically updates based on the key input. Ensure that the plot includes a descriptive title, a detailed legend reflecting tensor properties, and dynamic naming conventions based on the key string. The implementation should avoid hard coding and remain flexible, allowing any PyTorch module to be utilized seamlessly. Aim for a novel visualization approach that enhances interpretability and insight into the underlying data structure represented by the tensor.


directory folder save each tensor corresponding save .html .png ( automatically save) 



refence code 
import os
from typing import Any, Dict, List, Optional, Union
import logging
from pathlib import Path

import torch
import torch.nn as nn
import plotly.graph_objs as go
from plotly.subplots import make_subplots
import numpy as np

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class ModelVisualizer:
    """A class for visualizing tensors from various model types."""

    def __init__(self, output_dir: Union[str, Path] = "visualizations"):
        """
        Initialize the ModelVisualizer.

        Args:
            output_dir (Union[str, Path]): Directory to save visualization outputs.
        """
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        logger.info(f"Visualizations will be saved in: {self.output_dir}")

    def visualize_tensor(self, tensor: torch.Tensor, metadata: Dict[str, Any]) -> None:
        """
        Visualize a tensor based on its dimensionality.

        Args:
            tensor (torch.Tensor): The tensor to visualize.
            metadata (Dict[str, Any]): Additional information for visualization.
        """
        dim = tensor.dim()
        logger.info(f"Visualizing a tensor with {dim} dimensions.")

        try:
            if dim == 1:
                self._plot_1d(tensor, metadata)
            elif dim == 2:
                self._plot_2d(tensor, metadata)
            elif dim == 3:
                self._plot_3d(tensor, metadata)
            elif dim == 4:
                self._plot_4d(tensor, metadata)
            else:
                logger.warning(f"Tensors with {dim} dimensions are not supported for visualization.")
        except Exception as e:
            logger.error(f"Error in visualizing tensor: {e}")
            raise

    def _plot_1d(self, tensor: torch.Tensor, metadata: Dict[str, Any]) -> None:
        """Generate a 1D line plot."""
        x = torch.arange(tensor.size(0)).numpy()
        y = tensor.numpy()
        
        fig = go.Figure(data=go.Scatter(x=x, y=y, mode='lines+markers', name='1D Plot'))
        title = self._generate_title(metadata, "1D Line Plot")
        fig.update_layout(
            title=title,
            xaxis_title='Index',
            yaxis_title='Value',
            template='plotly_white'
        )
        self._save_and_show(fig, "1d_plot")

    def _plot_2d(self, tensor: torch.Tensor, metadata: Dict[str, Any]) -> None:
        """Generate a 2D heatmap."""
        fig = go.Figure(data=go.Heatmap(z=tensor.numpy(), colorscale='Viridis'))
        title = self._generate_title(metadata, "2D Heatmap")
        fig.update_layout(
            title=title,
            xaxis_title='X Axis',
            yaxis_title='Y Axis',
            template='plotly_white'
        )
        self._save_and_show(fig, "2d_heatmap")

    def _plot_3d(self, tensor: torch.Tensor, metadata: Dict[str, Any]) -> None:
        """Generate a 3D surface plot."""
        x = y = np.arange(tensor.shape[0])
        X, Y = np.meshgrid(x, y)
        
        fig = go.Figure(data=[go.Surface(z=tensor.numpy(), x=X, y=Y)])
        title = self._generate_title(metadata, "3D Surface Plot")
        fig.update_layout(
            title=title,
            scene=dict(
                xaxis_title='X Axis',
                yaxis_title='Y Axis',
                zaxis_title='Z Axis'
            ),
            template='plotly_white'
        )
        self._save_and_show(fig, "3d_surface")

    def _plot_4d(self, tensor: torch.Tensor, metadata: Dict[str, Any]) -> None:
        """Generate multiple 3D plots for 4D tensor."""
        fig = make_subplots(
            rows=2, cols=2,
            specs=[[{'type': 'surface'}, {'type': 'surface'}],
                   [{'type': 'surface'}, {'type': 'surface'}]],
            subplot_titles=[f"Slice {i}" for i in range(4)]
        )

        for i in range(4):
            X, Y = np.meshgrid(range(tensor.shape[1]), range(tensor.shape[2]))
            fig.add_trace(
                go.Surface(z=tensor[i].numpy(), x=X, y=Y, showscale=False),
                row=(i // 2) + 1, col=(i % 2) + 1
            )

        title = self._generate_title(metadata, "4D Tensor Visualization")
        fig.update_layout(
            title=title,
            height=800,
            width=800,
            template='plotly_white'
        )
        self._save_and_show(fig, "4d_multi_surface")

    def _generate_title(self, metadata: Dict[str, Any], plot_type: str) -> str:
        """Generate a descriptive title based on metadata."""
        model_type = metadata.get('model_type', 'Unknown')
        layer_name = metadata.get('layer_name', 'Unknown')
        return f"{plot_type} - {model_type} Model, Layer: {layer_name}"

    def _save_and_show(self, fig: go.Figure, plot_name: str) -> None:
        """Save the figure as HTML and display it."""
        file_path = self.output_dir / f"{plot_name}.html"
        fig.write_html(str(file_path))
        logger.info(f"Plot saved to {file_path}")
        fig.show()


class AdvancedModelWrapper(nn.Module):
    """A wrapper class for advanced model visualization."""

    def __init__(self, model: nn.Module, visualizer: ModelVisualizer):
        """
        Initialize the AdvancedModelWrapper.

        Args:
            model (nn.Module): The PyTorch model to wrap.
            visualizer (ModelVisualizer): The visualizer instance.
        """
        super().__init__()
        self.model = model
        self.visualizer = visualizer
        self.hooks = []
        self._register_hooks()

    def _register_hooks(self) -> None:
        """Register forward hooks for all modules in the model."""
        for name, module in self.model.named_modules():
            if isinstance(module, (nn.Conv2d, nn.Linear, nn.LSTM, nn.Transformer)):
                self.hooks.append(
                    module.register_forward_hook(self._make_hook(name))
                )

    def _make_hook(self, name: str):
        """Create a hook function for the given module name."""
        def hook(module, input, output):
            metadata = {
                'model_type': type(self.model).__name__,
                'layer_name': name,
                'layer_type': type(module).__name__
            }
            self.visualizer.visualize_tensor(output, metadata)
        return hook
   
    def forward(self, *args, **kwargs):
        """
        Forward pass through the wrapped model.

        Args:
            *args: Variable length argument list.
            **kwargs: Arbitrary keyword arguments.

        Returns:
            The output of the wrapped model.
        """
        try:
            return self.model(*args, **kwargs)
        except Exception as e:
            logger.error(f"Error in model forward pass: {e}")
            raise

    def remove_hooks(self):
        """Remove all registered hooks."""
        for hook in self.hooks:
            hook.remove()
        self.hooks.clear()
        logger.info("All hooks have been removed.")


def visualize_model(model: nn.Module, input_data: torch.Tensor, output_dir: str = "visualizations") -> None:
    """
    Visualize the internals of a PyTorch model.

    Args:
        model (nn.Module): The PyTorch model to visualize.
        input_data (torch.Tensor): Sample input data for the model.
        output_dir (str): Directory to save visualization outputs.
    """
    try:
        visualizer = ModelVisualizer(output_dir)
        wrapped_model = AdvancedModelWrapper(model, visualizer)
        
        logger.info("Starting model visualization...")
        with torch.no_grad():
            wrapped_model(input_data)
        
        wrapped_model.remove_hooks()
        logger.info("Model visualization completed successfully.")
    except Exception as e:
        logger.error(f"Error in model visualization: {e}")
        raise