in the below code make generalised and append all below functionalities and arguments using concept *, ** to all access to user 


Convolution Layers

Pooling layers

Padding Layers
Normalization Layers
Recurrent Layers
Transformer Layers
Linear Layers
Dropout Layers
Sparse Layers
import copy
import logging
from typing import Optional, Type, Union, Callable, Any

import torch
import torch.nn as nn
from torch import Tensor
from rich.console import Console
from rich.table import Table
from rich.tree import Tree
from rich.panel import Panel

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class QuantizationError(Exception):
    """Custom exception for quantization errors."""
    pass


def round_clip(x: Tensor, a: float, b: float) -> Tensor:
    """
    Rounds the input tensor and clips its values to be within [a, b].

    Args:
        x (Tensor): Input tensor.
        a (float): Minimum clipping value.
        b (float): Maximum clipping value.

    Returns:
        Tensor: Rounded and clipped tensor.
    """
    return torch.clamp(torch.round(x), min=a, max=b)


def compute_gamma(w: Tensor, epsilon: float = 1e-8) -> Tensor:
    """
    Computes the scaling factor gamma for the weight matrix.

    Args:
        w (Tensor): Weight tensor.
        epsilon (float, optional): Small constant to prevent division by zero. Defaults to 1e-8.

    Returns:
        Tensor: Scaling factor gamma.

    Raises:
        QuantizationError: If computation of gamma fails.
    """
    try:
        gamma = torch.mean(torch.abs(w)) + epsilon
        return gamma
    except Exception as e:
        logger.error(f"Error computing gamma: {e}")
        raise QuantizationError(f"Error computing gamma: {e}")


def bitnet_b158_quantize_weights(w: Tensor, epsilon: float = 1e-8) -> Tensor:
    """
    Quantizes the weights using the BitNet b1.58 quantization method.

    Args:
        w (Tensor): Original weight tensor.
        epsilon (float, optional): Small constant to prevent division by zero. Defaults to 1e-8.

    Returns:
        Tensor: Quantized weight tensor.

    Raises:
        QuantizationError: If quantization fails.
    """
    gamma = compute_gamma(w, epsilon)
    scaled_w = w / gamma
    quantized_w = round_clip(scaled_w, -1.0, 1.0)
    return quantized_w * gamma


class WeightQuantizer:
    """Class to handle weight quantization."""

    def __init__(self, epsilon: float = 1e-8) -> None:
        """
        Initializes the WeightQuantizer.

        Args:
            epsilon (float, optional): Small constant to prevent division by zero. Defaults to 1e-8.
        """
        self.epsilon = epsilon

    def quantize(self, w: Tensor) -> Tensor:
        """
        Quantizes the input weight tensor.

        Args:
            w (Tensor): Weight tensor to quantize.

        Returns:
            Tensor: Quantized weight tensor.
        """
        return bitnet_b158_quantize_weights(w, self.epsilon)


class ActivationQuantizer:
    """Class to handle activation quantization to 8 bits."""

    def __init__(self, num_bits: int = 8) -> None:
        """
        Initializes the ActivationQuantizer.

        Args:
            num_bits (int, optional): Number of bits for quantization. Defaults to 8.

        Raises:
            ValueError: If num_bits is not within [1, 16].
        """
        if not 1 <= num_bits <= 16:
            raise ValueError("num_bits must be between 1 and 16.")
        self.num_bits = num_bits

    def quantize(self, x: Tensor) -> Tensor:
        """
        Quantizes the input activation tensor to the specified number of bits.

        Args:
            x (Tensor): Activation tensor to quantize.

        Returns:
            Tensor: Quantized activation tensor.

        Raises:
            QuantizationError: If quantization fails.
        """
        try:
            qmin = 0
            qmax = 2 ** self.num_bits - 1
            min_val = x.min()
            max_val = x.max()
            scale = (max_val - min_val) / (qmax - qmin)
            scale = max(scale, 1e-8)  # Prevent division by zero
            zero_point = qmin - torch.round(min_val / scale)
            zero_point = zero_point.clamp(qmin, qmax).int()
            x_quant = torch.clamp(torch.round(x / scale) + zero_point, qmin, qmax)
            x_dequant = (x_quant - zero_point) * scale
            return x_dequant
        except Exception as e:
            logger.error(f"Error during activation quantization: {e}")
            raise QuantizationError(f"Error during activation quantization: {e}")


class QuantizedModule(nn.Module):
    """
    Generic wrapper for quantized modules (e.g., Linear, Conv2d).
    """

    def __init__(
        self,
        module: nn.Module,
        weight_quantizer: WeightQuantizer,
        activation_quantizer: ActivationQuantizer,
        module_type: Type[nn.Module]
    ) -> None:
        """
        Initializes the QuantizedModule.

        Args:
            module (nn.Module): Original PyTorch module to quantize.
            weight_quantizer (WeightQuantizer): Quantizer for the weights.
            activation_quantizer (ActivationQuantizer): Quantizer for the activations.
            module_type (Type[nn.Module]): The type of the module (e.g., nn.Linear, nn.Conv2d).
        """
        super(QuantizedModule, self).__init__()
        self.module_type = module_type
        self.activation_quantizer = activation_quantizer

        # Initialize the quantized module with quantized weights
        if isinstance(module, nn.Linear):
            quantized_weight = weight_quantizer.quantize(module.weight.data.clone())
            self.module = module_type(
                module.in_features,
                module.out_features,
                bias=module.bias is not None
            )
            self.module.weight = nn.Parameter(quantized_weight)
            if module.bias is not None:
                self.module.bias = nn.Parameter(module.bias.data.clone())

        elif isinstance(module, nn.Conv2d):
            quantized_weight = weight_quantizer.quantize(module.weight.data.clone())
            self.module = module_type(
                module.in_channels,
                module.out_channels,
                kernel_size=module.kernel_size,
                stride=module.stride,
                padding=module.padding,
                dilation=module.dilation,
                groups=module.groups,
                bias=module.bias is not None,
                padding_mode=module.padding_mode
            )
            self.module.weight = nn.Parameter(quantized_weight)
            if module.bias is not None:
                self.module.bias = nn.Parameter(module.bias.data.clone())

        else:
            raise TypeError(f"Unsupported module type: {type(module)}")

    def forward(self, input: Tensor) -> Tensor:
        """
        Forward pass with quantized weights and activations.

        Args:
            input (Tensor): Input tensor.

        Returns:
            Tensor: Output tensor after module operation and activation quantization.
        """
        try:
            output = self.module(input)
            output = self.activation_quantizer.quantize(output)
            return output
        except Exception as e:
            logger.error(f"Error in QuantizedModule forward pass: {e}")
            raise QuantizationError(f"Error in QuantizedModule forward pass: {e}")


def replace_modules(
    model: nn.Module,
    module_types: Union[Type[nn.Module], tuple],
    quantized_module_factory: Callable[[nn.Module], nn.Module]
) -> nn.Module:
    """
    Recursively replaces specified module types in the model with quantized modules.

    Args:
        model (nn.Module): Original PyTorch model.
        module_types (Union[Type[nn.Module], tuple]): Module types to replace (e.g., nn.Linear).
        quantized_module_factory (Callable[[nn.Module], nn.Module]): Factory function to create quantized modules.

    Returns:
        nn.Module: Model with specified modules replaced by quantized modules.
    """
    for name, module in model.named_children():
        if isinstance(module, module_types):
            setattr(model, name, quantized_module_factory(module))
        else:
            replace_modules(module, module_types, quantized_module_factory)
    return model


class BitNetPipeline:
    """Class to handle the full BitNet b1.58 quantization pipeline."""

    def __init__(self, model: nn.Module) -> None:
        """
        Initializes the BitNetPipeline.

        Args:
            model (nn.Module): The PyTorch model to quantize.
        """
        self.original_model = model
        self.quantized_model = copy.deepcopy(model)
        self.weight_quantizer = WeightQuantizer()
        self.activation_quantizer = ActivationQuantizer()

    def quantize_module(self, module: nn.Module) -> nn.Module:
        """
        Factory method to create a quantized version of the given module.

        Args:
            module (nn.Module): The module to quantize.

        Returns:
            nn.Module: Quantized module.
        """
        if isinstance(module, nn.Linear):
            return QuantizedModule(
                module,
                self.weight_quantizer,
                self.activation_quantizer,
                nn.Linear
            )
        elif isinstance(module, nn.Conv2d):
            return QuantizedModule(
                module,
                self.weight_quantizer,
                self.activation_quantizer,
                nn.Conv2d
            )
        else:
            raise TypeError(f"Unsupported module type for quantization: {type(module)}")

    def quantize(self) -> nn.Module:
        """
        Applies the quantization pipeline to the model.

        Returns:
            nn.Module: Quantized PyTorch model.

        Raises:
            QuantizationError: If quantization fails.
        """
        try:
            self.quantized_model = replace_modules(
                self.quantized_model,
                (nn.Linear, nn.Conv2d),
                self.quantize_module
            )
            logger.info("Model quantization completed successfully.")
            return self.quantized_model
        except Exception as e:
            logger.error(f"Error during model quantization: {e}")
            raise QuantizationError(f"Error during model quantization: {e}")


def apply_bitnet_b158_quantization(model: nn.Module) -> nn.Module:
    """
    Applies the BitNet b1.58 quantization pipeline to any PyTorch model.

    Args:
        model (nn.Module): The PyTorch model to quantize.

    Returns:
        nn.Module: Quantized PyTorch model.

    Raises:
        TypeError: If the input is not a PyTorch nn.Module.
        QuantizationError: If quantization fails.
    """
    if not isinstance(model, nn.Module):
        raise TypeError("Input model must be an instance of torch.nn.Module.")

    pipeline = BitNetPipeline(model)
    quantized_model = pipeline.quantize()
    return quantized_model


def model_summary(model: nn.Module) -> None:
    """
    Generate a detailed summary of a PyTorch model, including the number of total,
    trainable, and non-trainable parameters. Also, visually display the model architecture
    using a tree structure.

    Args:
        model (nn.Module): The PyTorch model to summarize.

    Raises:
        TypeError: If the input is not a PyTorch nn.Module.
    """
    if not isinstance(model, nn.Module):
        raise TypeError("Input must be a PyTorch nn.Module")

    console = Console()

    # Compute parameters' stats
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    non_trainable_params = total_params - trainable_params

    # Create a caption for the model
    model_caption = f"[bold cyan]{model.__class__.__name__}[/bold cyan]"
    console.print(Panel(model_caption, expand=False))

    # Create a summary table for the total parameters
    total_params_table = Table(
        title="Total Parameters Summary",
        show_header=True,
        header_style="bold magenta"
    )
    total_params_table.add_column("Category", style="cyan")
    total_params_table.add_column("Count", style="green", justify="right")

    total_params_table.add_row("Total Parameters", f"{total_params:,}")
    total_params_table.add_row("Trainable Parameters", f"{trainable_params:,}")
    total_params_table.add_row("Non-trainable Parameters", f"{non_trainable_params:,}")

    console.print(total_params_table)
    console.print(Panel("[italic]Summary of model's total parameters[/italic]", expand=False))

    # Create a detailed table for the parameters
    detailed_table = Table(
        title="Detailed Parameter Summary",
        show_lines=True
    )
    detailed_table.add_column("Parameter", style="cyan", no_wrap=True)
    detailed_table.add_column("Shape", style="magenta")
    detailed_table.add_column("Trainable", style="green")
    detailed_table.add_column("Num Parameters", style="yellow", justify="right")

    def add_module_to_table(module: nn.Module, prefix: Optional[str] = "") -> None:
        """
        Recursively add model parameters to the detailed table.

        Args:
            module (nn.Module): A PyTorch module.
            prefix (Optional[str]): Prefix for parameter names, used for recursion.
        """
        for name, param in module.named_parameters(recurse=False):
            full_name = f"{prefix}.{name}" if prefix else name
            shape = tuple(param.shape)
            trainable = param.requires_grad
            num_params = param.numel()

            detailed_table.add_row(
                full_name,
                str(shape),
                "[green]✓[/green]" if trainable else "[red]✗[/red]",
                f"{num_params:,}"
            )

        for name, child in module.named_children():
            child_prefix = f"{prefix}.{name}" if prefix else name
            add_module_to_table(child, child_prefix)

    add_module_to_table(model)
    console.print(detailed_table)
    console.print(Panel("[italic]Detailed breakdown of model's parameters[/italic]", expand=False))

    # Create a rich Tree for the model architecture
    def build_tree(module: nn.Module, prefix: Optional[str] = "") -> Tree:
        """
        Recursively build a tree structure for the model architecture.

        Args:
            module (nn.Module): The module to represent in the tree.
            prefix (Optional[str]): The prefix for the current module.

        Returns:
            Tree: A rich Tree representing the module architecture.
        """
        root_name = prefix if prefix else module.__class__.__name__
        tree = Tree(f"[bold]{root_name}[/bold]")

        for name, child in module.named_children():
            child_tree = build_tree(child, prefix=name)
            tree.add(child_tree)

        for name, param in module.named_parameters(recurse=False):
            param_info = (
                f"{name}: {tuple(param.shape)} "
                f"{'(Trainable)' if param.requires_grad else '(Non-trainable)'}"
            )
            tree.add(f"[cyan]{param_info}[/cyan]")

        return tree

    architecture_tree = build_tree(model)
    console.print("\n[bold yellow]Model Architecture[/bold yellow]")
    console.print(architecture_tree)
    console.print(Panel("[italic]Visual representation of the model's architecture[/italic]", expand=False))


def model_summary_update(model: nn.Module, save_path: Optional[str] = "model_summary.html") -> None:
    """
    Generate a detailed summary of a PyTorch model, including the number of total,
    trainable, and non-trainable parameters. Also, visually display the model architecture
    using a tree structure. Optionally saves the summary as an HTML file.

    Args:
        model (nn.Module): The PyTorch model to summarize.
        save_path (Optional[str]): Path to save the summary as an HTML file. Defaults to "model_summary.html".

    Raises:
        TypeError: If the input is not a PyTorch nn.Module.
    """
    if not isinstance(model, nn.Module):
        raise TypeError("Input must be a PyTorch nn.Module")

    console = Console(record=True)  # Enable recording of console output

    # Compute parameters' stats
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    non_trainable_params = total_params - trainable_params

    # Create a caption for the model
    model_caption = f"[bold cyan]{model.__class__.__name__}[/bold cyan]"
    console.print(Panel(model_caption, expand=False))

    # Create a summary table for the total parameters
    total_params_table = Table(
        title="Total Parameters Summary",
        show_header=True,
        header_style="bold magenta"
    )
    total_params_table.add_column("Category", style="cyan")
    total_params_table.add_column("Count", style="green", justify="right")

    total_params_table.add_row("Total Parameters", f"{total_params:,}")
    total_params_table.add_row("Trainable Parameters", f"{trainable_params:,}")
    total_params_table.add_row("Non-trainable Parameters", f"{non_trainable_params:,}")

    console.print(total_params_table)
    console.print(Panel("[italic]Summary of model's total parameters[/italic]", expand=False))

    # Create a detailed table for the parameters
    detailed_table = Table(
        title="Detailed Parameter Summary",
        show_lines=True
    )
    detailed_table.add_column("Parameter", style="cyan", no_wrap=True)
    detailed_table.add_column("Shape", style="magenta")
    detailed_table.add_column("Trainable", style="green")
    detailed_table.add_column("Num Parameters", style="yellow", justify="right")

    def add_module_to_table(module: nn.Module, prefix: Optional[str] = "") -> None:
        """
        Recursively add model parameters to the detailed table.

        Args:
            module (nn.Module): A PyTorch module.
            prefix (Optional[str]): Prefix for parameter names, used for recursion.
        """
        for name, param in module.named_parameters(recurse=False):
            full_name = f"{prefix}.{name}" if prefix else name
            shape = tuple(param.shape)
            trainable = param.requires_grad
            num_params = param.numel()

            detailed_table.add_row(
                full_name,
                str(shape),
                "[green]✓[/green]" if trainable else "[red]✗[/red]",
                f"{num_params:,}"
            )

        for name, child in module.named_children():
            child_prefix = f"{prefix}.{name}" if prefix else name
            add_module_to_table(child, child_prefix)

    add_module_to_table(model)
    console.print(detailed_table)
    console.print(Panel("[italic]Detailed breakdown of model's parameters[/italic]", expand=False))

    # Create a rich Tree for the model architecture
    def build_tree(module: nn.Module, prefix: Optional[str] = "") -> Tree:
        """
        Recursively build a tree structure for the model architecture.

        Args:
            module (nn.Module): The module to represent in the tree.
            prefix (Optional[str]): The prefix for the current module.

        Returns:
            Tree: A rich Tree representing the module architecture.
        """
        root_name = prefix if prefix else module.__class__.__name__
        tree = Tree(f"[bold]{root_name}[/bold]")

        for name, child in module.named_children():
            child_tree = build_tree(child, prefix=name)
            tree.add(child_tree)

        for name, param in module.named_parameters(recurse=False):
            param_info = (
                f"{name}: {tuple(param.shape)} "
                f"{'(Trainable)' if param.requires_grad else '(Non-trainable)'}"
            )
            tree.add(f"[cyan]{param_info}[/cyan]")

        return tree

    architecture_tree = build_tree(model)
    console.print("\n[bold yellow]Model Architecture[/bold yellow]")
    console.print(architecture_tree)
    console.print(Panel("[italic]Visual representation of the model's architecture[/italic]", expand=False))

    if save_path:
        try:
            html = console.export_html(clear=False)
            with open(save_path, "w", encoding="utf-8") as f:
                f.write(html)
            console.print(f"Summary saved as [green]{save_path}[/green]")
        except Exception as e:
            logger.error(f"Failed to save summary to {save_path}: {e}")
            raise QuantizationError(f"Failed to save summary to {save_path}: {e}")