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 :user model (torch.nn module) generalised pipeline any model from pytorch module , we apply below algorithm,note activation 8 bit , remaining all 1.5bit follow algorithm

## BitNet b1.58: A Detailed Explanation

BitNet b1.58 is a variant of the BitNet architecture, a type of Transformer model designed for Large Language Models (LLMs). This variant uses a ternary system where each parameter (weight) can be -1, 0, or 1, as opposed to the standard 16-bit floating values (FP16 or BF16) used in conventional LLMs. This results in significant cost reductions in terms of:

*   Latency
*   Memory usage
*   Throughput
*   Energy consumption

while maintaining comparable performance to full-precision models.

### Advantages of BitNet b1.58

*   **Pareto Improvement**: BitNet b1.58 provides a Pareto improvement by reducing inference costs while maintaining model performance. This means it achieves a better balance between cost and performance compared to traditional LLMs.
*   **Simplified Computation**: BitNet b1.58 replaces floating-point operations with integer additions for matrix multiplication, leading to significant energy savings and faster computation.
*   **Reduced Memory Footprint**: It has a lower memory footprint, enabling faster and more efficient loading of weights, and making it suitable for deployment on resource-constrained devices.
*   **Enhanced Modeling Capability**: The introduction of '0' in the weight values allows for explicit feature filtering, improving the model's ability to capture important information.
*   **Open-Source Compatibility**: BitNet b1.58 incorporates LLaMA-alike components (RMSNorm, SwiGLU, rotary embedding), ensuring easy integration with popular open-source LLM frameworks.

### Quantization Function

The key to achieving 1.58-bit precision lies in the quantization function used to constrain the weights to the ternary set {-1, 0, 1}. BitNet b1.58 uses the following quantization function:

**Equation 1:**

$ \tilde{W} = RoundClip(\frac{W}{γ + ϵ} ,−1, 1) $

**Equation 2:**

$RoundClip(x, a, b) = max(a,min(b, round(x))) $

**Equation 3:**

$ γ =  \frac{1}{nm}\sum_{ij} |W_{ij}| $

**Explanation:**

1.  The weight matrix (*W*) is scaled by its average absolute value (*γ*) plus a small constant (*ϵ*).
2.  The *RoundClip* function then rounds each scaled value to the nearest integer within the range of -1 to 1.

NOTE: rounding is not working below function update logic correct in the code


code:
import copy
import logging
from typing import Optional, Type, Union, Callable, Any, Tuple

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 a specified number of 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.
    Handles any nn.Module by quantizing weights (if present) and activations.
    """

    def __init__(
        self,
        module: nn.Module,
        weight_quantizer: WeightQuantizer,
        activation_quantizer: ActivationQuantizer
    ) -> 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.
        """
        super(QuantizedModule, self).__init__()
        self.activation_quantizer = activation_quantizer
        self.module = copy.deepcopy(module)  # Deep copy to preserve original module

        # Check if the module has weight parameters and quantize them
        if hasattr(self.module, 'weight') and isinstance(self.module.weight, nn.Parameter):
            with torch.no_grad():
                quantized_weight = weight_quantizer.quantize(self.module.weight.data.clone())
                self.module.weight = nn.Parameter(quantized_weight)
                logger.debug(f"Quantized weights for {self.module.__class__.__name__}")

    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[Type[nn.Module], ...]],
    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):
            logger.debug(f"Replacing module: {name} ({type(module).__name__})")
            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.
        """
        return QuantizedModule(
            module,
            self.weight_quantizer,
            self.activation_quantizer
        )

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

        Returns:
            nn.Module: Quantized PyTorch model.

        Raises:
            QuantizationError: If quantization fails.
        """
        try:
            # Define all module types that potentially have weights
            modules_with_weights = (
                nn.Conv1d,
                nn.Conv2d,
                nn.Conv3d,
                nn.Linear,
                nn.Embedding,
                nn.LayerNorm,
                nn.BatchNorm1d,
                nn.BatchNorm2d,
                nn.BatchNorm3d,
                nn.LSTM,
                nn.GRU,
                nn.MultiheadAttention,
                nn.RNN,
                nn.TransformerEncoderLayer,
                nn.TransformerDecoderLayer,
                nn.Parameter
                
                # Add more modules as needed
            )

            self.quantized_model = replace_modules(
                self.quantized_model,
                modules_with_weights,
                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}")