Multiprocessing failed: _plot_and_save() missing 1 required positional argument: 'tensor'. Falling back to single-threaded processing
import os
import multiprocessing
from typing import Any, Optional, Union
import torch
from torch import nn
import plotly.graph_objs as go
from plotly.subplots import make_subplots
from functools import partial


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 _plot_and_save(
    name: str,
    tensor: torch.Tensor,
    output_dir: str,
    max_dim: int = 2
) -> None:
    """
    Plots a single tensor and saves it as HTML and PNG files.

    Args:
        name (str): The name of the tensor.
        tensor (torch.Tensor): The tensor to plot.
        output_dir (str): The directory where plots will be saved.
        max_dim (int): Maximum dimensions to plot directly. Higher dimensions will be flattened.
    """
    try:
        tensor_cpu = tensor.cpu().detach()
        tensor_dim = tensor_cpu.dim()

        if tensor_dim == 0:
            # Scalar tensor
            fig = go.Figure()
            fig.add_trace(
                go.Scatter(
                    x=[0],
                    y=[tensor_cpu.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_cpu = _flatten_tensor(tensor_cpu)
                current_dim = 2
            else:
                current_dim = tensor_dim

            if current_dim == 1:
                x = list(range(tensor_cpu.size(0)))
                y = tensor_cpu.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_cpu.size(0)
                for i in range(num_rows):
                    x = list(range(tensor_cpu.size(1)))
                    y = tensor_cpu[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_cpu.size(0)):
                    x = list(range(tensor_cpu.size(1)))
                    y = tensor_cpu[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}")


def save_model_tensor_plots(
    model: nn.Module,
    output_dir: str,
    max_dim: int = 2,
    use_multiprocessing: bool = True
) -> 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.
        use_multiprocessing (bool): Whether to use multiprocessing for plotting.

    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)

    tensors = list(model.named_parameters())

    if use_multiprocessing:
        try:
            cpu_count = multiprocessing.cpu_count()
            with multiprocessing.Pool(processes=cpu_count) as pool:
                pool.map(
                    partial(
                        _plot_and_save,
                        output_dir=output_dir,
                        max_dim=max_dim
                    ),
                    [(name, param.data) for name, param in tensors]
                )
        except Exception as e:
            print(f"Multiprocessing failed: {e}. Falling back to single-threaded processing.")
            for name, param in tensors:
                _plot_and_save(name, tensor=param.data, output_dir=output_dir, max_dim=max_dim)
    else:
        for name, param in tensors:
            _plot_and_save(name, tensor=param.data, output_dir=output_dir, max_dim=max_dim)


# Example Usage
if __name__ == "__main__":
    try:
        from transformers import AutoModelForCausalLM
    except ImportError as e:
        raise ImportError("Please install the 'transformers' library to proceed.") from e

    try:
        model = AutoModelForCausalLM.from_pretrained("gpt2")
    except Exception as e:
        raise RuntimeError(f"Failed to load the model: {e}") from e

    # Define output directory
    output_directory = "model_tensor_plots"

    # Save tensor plots using all CPU cores
    save_model_tensor_plots(model, output_directory)