Diff Coverage

def _normalize_mhc_pre_clamp_sinkhorn_args(*args, **kwargs):
    """Normalize npu_mhc_pre_clamp_sinkhorn arguments."""
    values = dict(_MHC_PRE_CLAMP_DEFAULTS)
    if len(args) > len(_MHC_PRE_CLAMP_ARG_NAMES):
        raise TypeError(
            f"npu_mhc_pre_clamp_sinkhorn expected at most {len(_MHC_PRE_CLAMP_ARG_NAMES)} arguments"
        )
    for name, value in zip(_MHC_PRE_CLAMP_ARG_NAMES, args):
        values[name] = value
    for name, value in kwargs.items():
        if name not in _MHC_PRE_CLAMP_ARG_NAMES:
            raise TypeError(f"npu_mhc_pre_clamp_sinkhorn got an unexpected keyword argument '{name}'")
        if name in _MHC_PRE_CLAMP_ARG_NAMES[:len(args)]:
            raise TypeError(f"npu_mhc_pre_clamp_sinkhorn got multiple values for argument '{name}'")
        values[name] = value
    missing = [name for name in _MHC_PRE_CLAMP_ARG_NAMES[:4] if name not in values]
    if missing:
        raise TypeError(f"npu_mhc_pre_clamp_sinkhorn missing required arguments: {missing}")
    return tuple(values[name] for name in _MHC_PRE_CLAMP_ARG_NAMES), {}


# Validation rules table for npu_mhc_pre_sinkhorn
# Key: tensor_map length (format identifier)

        ]
        return local_args, local_kwargs, cache_values

    def infer_layout(self, layouts: list, extra_args=None) -> Tuple[tuple, None]:
        del extra_args
        x_layout, phi_layout, alpha_layout, bias_layout = layouts

        self._check_partial_inputs([x_layout, phi_layout, alpha_layout, bias_layout])
        _validate_input_layouts_mhc_pre_sinkhorn(
            x_layout, phi_layout, alpha_layout, bias_layout

def _bind_mhc_pre_clamp_args(args, kwargs):
    """Bind npu_mhc_pre_clamp_sinkhorn arguments with Python defaults."""
    names = (
        "x", "phi", "alpha", "bias", "hc_mult", "num_iters",
        "hc_eps", "norm_eps", "out_flag", "clamp_min", "clamp_max",
    )
    values = {
        "hc_mult": 4,
        "num_iters": 20,
        "hc_eps": 1e-6,
        "norm_eps": 1e-6,

        "out_flag": True,
        "clamp_min": 0.0,
        "clamp_max": 0.0,
    }
    if len(args) > len(names):
        raise TypeError(f"npu_mhc_pre_clamp_sinkhorn expected at most {len(names)} arguments")
    for name, value in zip(names, args):
        values[name] = value
    for name, value in kwargs.items():
        if name in values and name in names[:len(args)]:
            raise TypeError(f"npu_mhc_pre_clamp_sinkhorn got multiple values for argument '{name}'")
        if name not in names:
            raise TypeError(f"npu_mhc_pre_clamp_sinkhorn got an unexpected keyword argument '{name}'")
        values[name] = value
    missing = [name for name in names[:4] if name not in values]
    if missing:
        raise TypeError(f"npu_mhc_pre_clamp_sinkhorn missing required arguments: {missing}")
    return _MhcPreClampArgs(*(values[name] for name in names))


def _build_custom_ops():
    return ms.ops.CustomOpBuilder(

    @staticmethod
    def forward(ctx, *args, **kwargs):
        """Forward pass: delegates to the clamp-enabled Ascend custom kernel."""
        bound = _bind_mhc_pre_clamp_args(args, kwargs)
        result = _custom_ops.npu_mhc_pre_clamp_sinkhorn(
            bound.x, bound.phi, bound.alpha, bound.bias,
            bound.hc_mult, bound.num_iters, bound.hc_eps, bound.norm_eps,
            bound.out_flag, bound.clamp_min, bound.clamp_max

            bound.hc_mult, bound.num_iters, bound.hc_eps, bound.norm_eps,
            bound.out_flag, bound.clamp_min, bound.clamp_max
        )
        _, _, _, h_pre, hc_before_norm, inv_rms, sum_out, norm_out, h_res_logits = result
        ctx.save_for_backward(bound.x, bound.phi, bound.alpha, bound.bias,
                              h_pre, hc_before_norm, inv_rms, sum_out, norm_out, h_res_logits)
        ctx.hc_eps = bound.hc_eps
        ctx.clamp_min = bound.clamp_min
        ctx.clamp_max = bound.clamp_max
        return result

    @staticmethod
    def backward(ctx, *grad_outputs):

    @staticmethod
    def backward(ctx, *grad_outputs):
        """Backward pass: calls npu_mhc_pre_clamp_sinkhorn_backward kernel."""
        tensors = _ensure_contiguous(
            grad_outputs[0], grad_outputs[1], grad_outputs[2],
            *ctx.saved_tensors
        )
        n = tensors[1].shape[-1]
        grad_h_res = ms.ops.reshape(tensors[2], tuple(tensors[2].shape[:-1]) + (n, n))

        grads = _custom_ops.npu_mhc_pre_clamp_sinkhorn_backward(
            tensors[0], tensors[1], grad_h_res,
            tensors[3], tensors[4], tensors[5], tensors[6],

            tensors[0], tensors[1], grad_h_res,
            tensors[3], tensors[4], tensors[5], tensors[6],
            tensors[7], tensors[8], tensors[9], tensors[10], tensors[11], tensors[12],
            ctx.hc_eps, ctx.clamp_min, ctx.clamp_max)
        return tuple(grads[:4]) + _MHC_PRE_CLAMP_NONE_GRADS