TASK : each layer we visualization 
GPT2LMHeadModel(
  (transformer): GPT2Model(
    (wte): Embedding(50257, 768)
    (wpe): Embedding(1024, 768)
    (drop): Dropout(p=0.1, inplace=False)
    (h): ModuleList(
      (0-11): 12 x GPT2Block(
        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (attn): GPT2SdpaAttention(
          (c_attn): Conv1D()
          (c_proj): Conv1D()
          (attn_dropout): Dropout(p=0.1, inplace=False)
          (resid_dropout): Dropout(p=0.1, inplace=False)
        )
        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (mlp): GPT2MLP(
          (c_fc): Conv1D()
          (c_proj): Conv1D()
          (act): NewGELUActivation()
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
    )
    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
  )
  (lm_head): Linear(in_features=768, out_features=50257, bias=False)
)

import torch
import torch.nn as nn
import plotly.graph_objs as go
from plotly.subplots import make_subplots
from typing import Any, Dict, Union
import logging

# Configure logging for debugging and error tracking
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class DynamicTensorModel(nn.Module):
    """
    A PyTorch model that accepts a recursive key input and returns a tensor of varying dimensions.
    It also visualizes the tensor using Plotly based on its dimensionality.
    """

    def __init__(self) -> None:
        super(DynamicTensorModel, self).__init__()

    def forward(self, key_input: Dict[str, Any]) -> torch.Tensor:
        """
        Processes the recursive key input to generate a tensor.

        Args:
            key_input (Dict[str, Any]): A nested dictionary input.

        Returns:
            torch.Tensor: A tensor with dimensions based on the key input.
        """
        try:
            tensor = self.generate_tensor(key_input)
            self.visualize_tensor(tensor, key_input)
            return tensor
        except Exception as e:
            logger.error(f"Error in forward pass: {e}")
            raise

    def generate_tensor(self, key_input: Dict[str, Any], depth: int = 1) -> torch.Tensor:
        """
        Recursively processes the key input to generate a tensor with varying dimensions.

        Args:
            key_input (Dict[str, Any]): A nested dictionary input.
            depth (int, optional): Current depth of recursion. Defaults to 1.

        Returns:
            torch.Tensor: A tensor with dimensions based on recursion depth.
        """
        if not isinstance(key_input, dict):
            raise ValueError("Key input must be a dictionary.")

        if not key_input:
            raise ValueError("Key input dictionary is empty.")

        # Base case: when the value is not a dictionary
        for key, value in key_input.items():
            if isinstance(value, dict):
                return self.generate_tensor(value, depth + 1)
            else:
                # Create a tensor with dimension based on depth
                size = 10  # arbitrary size for each dimension
                dims = min(depth, 4)  # Limit to 4D
                shape = tuple([size] * dims)
                logger.info(f"Generating a {dims}D tensor with shape {shape} based on key '{key}'.")
                return torch.randn(shape)

        # If all values are dictionaries, increase depth
        size = 10
        dims = min(depth, 4)
        shape = tuple([size] * dims)
        logger.info(f"Generating a {dims}D tensor with shape {shape} based on nested keys.")
        return torch.randn(shape)

    def visualize_tensor(self, tensor: torch.Tensor, key_input: Dict[str, Any]) -> None:
        """
        Visualizes the tensor using Plotly based on its dimensionality.

        Args:
            tensor (torch.Tensor): The tensor to visualize.
            key_input (Dict[str, Any]): The original key input for dynamic naming.
        """
        dim = tensor.dim()
        logger.info(f"Visualizing a tensor with {dim} dimensions.")

        if dim == 1:
            self._plot_1d(tensor, key_input)
        elif dim == 2:
            self._plot_2d(tensor, key_input)
        elif dim == 3:
            self._plot_3d(tensor, key_input)
        elif dim == 4:
            self._plot_4d(tensor, key_input)
        else:
            logger.warning(f"Tensors with {dim} dimensions are not supported for visualization.")

    def _plot_1d(self, tensor: torch.Tensor, key_input: Dict[str, Any]) -> None:
        """Generates a simple 1D scatter plot."""
        try:
            x = torch.arange(tensor.size(0)).numpy()
            y = tensor.numpy()
            fig = go.Figure(data=go.Scatter(x=x, y=y, mode='markers', name='1D Scatter'))
            title = self._generate_title(key_input, "1D Scatter Plot")
            fig.update_layout(title=title, xaxis_title='Index', yaxis_title='Value')
            fig.show()
            logger.info("1D scatter plot generated successfully.")
        except Exception as e:
            logger.error(f"Error plotting 1D tensor: {e}")

    def _plot_2d(self, tensor: torch.Tensor, key_input: Dict[str, Any]) -> None:
        """Generates a standard 2D scatter plot."""
        try:
            if tensor.size(1) < 2:
                raise ValueError("2D tensor must have at least two features for plotting.")

            x = tensor[:, 0].numpy()
            y = tensor[:, 1].numpy()
            fig = go.Figure(data=go.Scatter(x=x, y=y, mode='markers', name='2D Scatter'))
            title = self._generate_title(key_input, "2D Scatter Plot")
            fig.update_layout(title=title, xaxis_title='Feature 1', yaxis_title='Feature 2')
            fig.show()
            logger.info("2D scatter plot generated successfully.")
        except Exception as e:
            logger.error(f"Error plotting 2D tensor: {e}")

    def _plot_3d(self, tensor: torch.Tensor, key_input: Dict[str, Any]) -> None:
        """Generates an interactive 3D scatter plot."""
        try:
            if tensor.size(1) < 3:
                raise ValueError("3D tensor must have at least three features for plotting.")

            x = tensor[:, 0].numpy()
            y = tensor[:, 1].numpy()
            z = tensor[:, 2].numpy()
            fig = go.Figure(data=go.Scatter3d(
                x=x, y=y, z=z,
                mode='markers',
                marker=dict(size=5),
                name='3D Scatter'
            ))
            title = self._generate_title(key_input, "3D Scatter Plot")
            fig.update_layout(title=title,
                              scene=dict(
                                  xaxis_title='Feature 1',
                                  yaxis_title='Feature 2',
                                  zaxis_title='Feature 3'
                              ))
            fig.show()
            logger.info("3D scatter plot generated successfully.")
        except Exception as e:
            logger.error(f"Error plotting 3D tensor: {e}")

    def _plot_4d(self, tensor: torch.Tensor, key_input: Dict[str, Any]) -> None:
        """Generates an advanced 3D scatter plot with a dynamic legend based on the key input."""
        try:
            if tensor.size(1) < 4:
                raise ValueError("4D tensor must have at least four features for plotting.")

            x = tensor[:, 0].numpy()
            y = tensor[:, 1].numpy()
            z = tensor[:, 2].numpy()
            color_feature = tensor[:, 3].numpy()

            # Create a color scale based on the fourth dimension
            colorscale = 'Viridis'

            fig = go.Figure(data=go.Scatter3d(
                x=x, y=y, z=z,
                mode='markers',
                marker=dict(
                    size=5,
                    color=color_feature,
                    colorscale=colorscale,
                    opacity=0.8,
                    colorbar=dict(title='Feature 4')
                ),
                name='4D Scatter'
            ))
            title = self._generate_title(key_input, "4D Scatter Plot with Dynamic Legend")
            fig.update_layout(title=title,
                              scene=dict(
                                  xaxis_title='Feature 1',
                                  yaxis_title='Feature 2',
                                  zaxis_title='Feature 3'
                              ))
            fig.show()
            logger.info("4D scatter plot generated successfully.")
        except Exception as e:
            logger.error(f"Error plotting 4D tensor: {e}")

    def _generate_title(self, key_input: Dict[str, Any], plot_type: str) -> str:
        """
        Generates a descriptive title based on the key input.

        Args:
            key_input (Dict[str, Any]): The original key input.
            plot_type (str): The type of plot.

        Returns:
            str: A descriptive title.
        """
        flat_keys = self._flatten_keys(key_input)
        key_str = "_".join(flat_keys)
        title = f"{plot_type} - Keys: {key_str}"
        return title

    def _flatten_keys(self, d: Dict[str, Any], parent_key: str = '', sep: str = '_') -> list:
        """
        Flattens nested dictionary keys.

        Args:
            d (Dict[str, Any]): The dictionary to flatten.
            parent_key (str, optional): The base key string. Defaults to ''.
            sep (str, optional): Separator between keys. Defaults to '_'.

        Returns:
            list: A list of flattened key strings.
        """
        items = []
        for k, v in d.items():
            new_key = f"{parent_key}{sep}{k}" if parent_key else k
            if isinstance(v, dict):
                items.extend(self._flatten_keys(v, new_key, sep=sep))
            else:
                items.append(new_key)
        return items