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 :

Two major drawback in our code is 
1) higher dim plot we are converting into flatting , we can use 3d, plots colour logic 
2) major is not making fast using cuda , cpu all mutiprocessing torch or module








import os
from typing import Any, Optional, Union
import torch
from torch import nn
import plotly.graph_objs as go
from plotly.subplots import make_subplots


def _flatten_tensor(tensor: torch.Tensor) -> torch.Tensor:
    """
    Flattens a tensor to 2D if it's higher than 2D.

    Args:
        tensor (torch.Tensor): The tensor to flatten.

    Returns:
        torch.Tensor: A 2D tensor.
    """
    if tensor.dim() == 1:
        return tensor.unsqueeze(1)
    elif tensor.dim() == 2:
        return tensor
    else:
        return tensor.view(tensor.size(0), -1)


def _ensure_directory(path: str) -> None:
    """
    Ensures that the given directory exists. Creates it if it does not.

    Args:
        path (str): The directory path to ensure.

    Raises:
        OSError: If the directory cannot be created.
    """
    try:
        os.makedirs(path, exist_ok=True)
    except OSError as e:
        raise OSError(f"Failed to create directory {path}: {e}") from e


def save_model_tensor_plots(
    model: nn.Module,
    output_dir: str,
    max_dim: int = 2
) -> None:
    """
    Recursively traverses the model's tensors, plots each tensor using Plotly,
    and saves the plots as .html and .png files in the specified directory.

    Args:
        model (nn.Module): The PyTorch model to traverse.
        output_dir (str): The directory where plots will be saved.
        max_dim (int): Maximum dimensions to plot directly. Higher dimensions will be flattened.

    Raises:
        TypeError: If model is not an instance of nn.Module.
        ValueError: If output_dir is not a valid directory path.
    """
    if not isinstance(model, nn.Module):
        raise TypeError("The model must be an instance of torch.nn.Module.")

    if not isinstance(output_dir, str) or not output_dir:
        raise ValueError("The output directory must be a non-empty string.")

    _ensure_directory(output_dir)

    for name, param in model.named_parameters():
        tensor = param.data.cpu().detach()
        tensor_dim = tensor.dim()

        try:
            if tensor_dim == 0:
                # Scalar tensor
                fig = go.Figure()
                fig.add_trace(go.Scatter(
                    x=[0],
                    y=[tensor.item()],
                    mode='markers',
                    name='Scalar Value'
                ))
                title = f"{name} (Scalar)"
                fig.update_layout(
                    title=title,
                    xaxis_title="Index",
                    yaxis_title="Value",
                    legend_title="Legend"
                )
            else:
                # Handle tensors with dimensions up to max_dim
                if tensor_dim > max_dim:
                    tensor = _flatten_tensor(tensor)
                    current_dim = 2
                else:
                    current_dim = tensor_dim

                if current_dim == 1:
                    x = list(range(tensor.size(0)))
                    y = tensor.tolist()
                    fig = go.Figure()
                    fig.add_trace(go.Scatter(
                        x=x,
                        y=y,
                        mode='markers',
                        name=name
                    ))
                    title = f"{name} (1D Tensor)"
                    fig.update_layout(
                        title=title,
                        xaxis_title="Index",
                        yaxis_title="Value",
                        legend_title="Legend"
                    )
                elif current_dim == 2:
                    # For 2D tensors, plot each row as a separate scatter
                    fig = go.Figure()
                    num_rows = tensor.size(0)
                    for i in range(num_rows):
                        x = list(range(tensor.size(1)))
                        y = tensor[i].tolist()
                        fig.add_trace(go.Scatter(
                            x=x,
                            y=y,
                            mode='markers',
                            name=f"{name} Row {i}"
                        ))
                    title = f"{name} (2D Tensor)"
                    fig.update_layout(
                        title=title,
                        xaxis_title="Dimension 1",
                        yaxis_title="Dimension 2",
                        legend_title="Legend"
                    )
                else:
                    # For higher dimensions (after flattening)
                    fig = make_subplots(rows=1, cols=1)
                    for i in range(tensor.size(0)):
                        x = list(range(tensor.size(1)))
                        y = tensor[i].tolist()
                        fig.add_trace(go.Scatter(
                            x=x,
                            y=y,
                            mode='markers',
                            name=f"{name} Slice {i}"
                        ))
                    title = f"{name} (Flattened Tensor)"
                    fig.update_layout(
                        title=title,
                        xaxis_title="Flattened Dimension 1",
                        yaxis_title="Flattened Dimension 2",
                        legend_title="Legend"
                    )

            # Define file paths
            safe_name = name.replace('.', '_').replace('/', '_')
            html_path = os.path.join(output_dir, f"{safe_name}.html")
            png_path = os.path.join(output_dir, f"{safe_name}.png")

            # Save as HTML
            fig.write_html(html_path)

            # Save as PNG
            fig.write_image(png_path)

        except Exception as e:
            print(f"Failed to plot tensor '{name}': {e}")


# Example Usage
if __name__ == "__main__":
    # Sample model for demonstration
    from transformers import AutoModelForCausalLM

    model=AutoModelForCausalLM.from_pretrained("gpt2")
    # Define output directory
    output_directory = "model_tensor_plots"
    # Save tensor plots
    save_model_tensor_plots(model, output_directory)