Diff Coverage

    """Flatten leading dimensions while preserving the contracting axis."""

    if tensor.ndim == 2:
        return tensor
    return tensor.reshape(-1, tensor.shape[-1])


class LinearStrategy(ABC):
    """Own the common Dense forward/backward flow and expose format hooks."""

            save_quantized_operands(ctx, *operands)
            ctx.input_quant = None
            ctx.weight_quant = None
        else:
            ctx.input_quant, ctx.weight_quant = operands
        if inputs.ndim != 2:
            output = output.reshape(*inputs.shape[:-1], output.shape[-1])
        return output

    def backward(
        self,

        needs_grad_weight = ctx.needs_input_grad[1]
        if ctx.quantized_operands_saved:
            _, input_quant, weight_quant = restore_quantized_operands(ctx)
        else:
            input_quant = ctx.input_quant
            weight_quant = ctx.weight_quant
        grad_quant = self.quantize_grad_output(
            grad_matrix,
            rowwise=needs_grad_input,
            colwise=needs_grad_weight,

                layout="NN",
                output_dtype=grad_output.dtype,
            )
            if len(ctx.input_shape) != 2:
                grad_input = grad_input.reshape(ctx.input_shape)
        if needs_grad_weight:
            grad_weight = self.matmul(
                grad_quant,
                input_quant,

        if inputs.dtype not in (torch.float16, torch.bfloat16) or weight.dtype not in (
            torch.float16,
            torch.bfloat16,
        ):
            raise TypeError(
                "HiFloat8 grouped linear inputs and weight must use float16 or bfloat16."
            )
        if group_list.dtype != torch.int64:
            raise TypeError("HiFloat8 grouped linear group_list must use torch.int64.")

        NotImplementedError: If the selected Dense policy is not implemented.
    """

    if isinstance(config, LowPrecisionConfig):
        config = config.resolve_dtype_scheme()
    if config is None:
        config = LowPrecisionDtypeScheme()

    if (
        not config.is_fake_quantize
        and config.weight_format == "mxfp8"

            )
        if weight_format == "hif8" and act_format == "hif8":
            validate_hifloat8_gmm_runtime()
            return HiFloat8GroupedLinear()
        raise NotImplementedError(
            "hifloat combination weight_format="
            f"{weight_format!r}/act_format={act_format!r} is not implemented; "
            "supported: hif8/hif8."
        )

        fqn: str = "",
    ) -> None:
        """Create an unbound packed-expert module."""

        self._validate_grouped_linear(grouped_linear)
        super().__init__()
        self.gate_up_proj = nn.Parameter(
            torch.empty(num_experts, 2 * intermediate_dim, hidden_dim)
        )

        self.num_experts = num_experts
        self.hidden_dim = hidden_dim
        self.intermediate_dim = intermediate_dim
        self.fqn = fqn
        self.grouped_linear = grouped_linear

    @staticmethod
    def _validate_grouped_linear(grouped_linear: GroupedLinear) -> None:
        """Require a plain strategy so module registration stays unchanged."""

        if isinstance(grouped_linear, nn.Module) or not isinstance(
            grouped_linear,
            GroupedLinear,
        ):
            raise TypeError(
                "grouped_linear must be a non-Module GroupedLinear instance"
            )

    @classmethod

        strategy: LinearStrategy,
    ) -> None:
        """Create a low-precision Linear with newly allocated Parameters."""

        self._validate_strategy(strategy)
        super().__init__(in_features, out_features, bias=bias)
        self.strategy = strategy

    @staticmethod
    def _validate_strategy(strategy: LinearStrategy) -> None:
        """Require a plain strategy so module registration stays unchanged."""

    def _validate_strategy(strategy: LinearStrategy) -> None:
        """Require a plain strategy so module registration stays unchanged."""

        if isinstance(strategy, nn.Module) or not isinstance(strategy, LinearStrategy):
            raise TypeError("strategy must be a non-Module LinearStrategy instance")

    @classmethod
    def from_linear(
        cls,

    def forward(self, input: torch.Tensor) -> torch.Tensor:  # pylint: disable=redefined-builtin
        """Apply the selected low-precision compute and high-precision bias."""

        output = _LinearFunction.apply(input, self.weight, self.strategy)
        if self.bias is not None:
            output = output + self.bias
        return output