import torch
from torch import Tensor

def quant_matmul(
    mat1: Tensor,
    mat2: Tensor,
    scale: Tensor,
    perTokenScale: Tensor,
    outDType: str | torch.dtype = torch.float16,
    transA: bool = False,
    transB: bool = False,
    useNzA: bool = False,
    useNzB: bool = False,
) -> Tensor:
    if isinstance(outDType, str):
        dtype_lower = outDType.lower()
        outDType = getattr(torch, dtype_lower, None)
    if outDType is None:
        raise ValueError(f"{outDType} is not a data type of torch")
    return torch.ops.catlass.quant_matmul(
        mat1, mat2, scale, perTokenScale, outDType, transA, transB, useNzA, useNzB
    )