Repository: pytorch
Issue: [ONNX] Exporting PyTorch Model to ONNX: Unsupported Operator 'aten::linalg_solve'
Labels: module: onnx,triaged,onnx-triaged

### 🐛 Describe the bug

I'm currently trying to export my PyTorch model to ONNX , but I'm encountering an error related to the unsupported operator aten::linalg_solve. Here is the error traceback:

```
2024-08-09 22:28:24 [ERROR] Error during the ONNX export and quantization process: Exporting the operator 'aten::linalg_solve' to ONNX opset version 14 is not supported. Please feel free to request sup
port or submit a pull request on PyTorch GitHub: https://github.com/pytorch/pytorch/issues.

```

However, my code does not explicitly use the linalg_solve function. Here's the relevant portion of my model code:

```
def eval_func(model):
    model.to(model_device)  # 确保评估时模型在正确的设备上


    model.eval()
    all_metrics = []

    loss_func = PITLossWrapper(pairwise_neg_snr, pit_from='pw_mtx')

    with torch.no_grad():
        for batch_idx, batch in enumerate(tqdm(val_loader1, desc="Evaluating")):  # 使用 tqdm 显示进度条
            mix, sources = tensors_to_device(batch, device=model_device)

            logger.debug(f"Processing batch {batch_idx}...")
            start_time = datetime.now()  # 记录开始时间

            # Forward the network on the mixture.
            forward_start_time = datetime.now()  # Record forward pass start time
            mix = mix.squeeze(0).to(model_device)  # Remove batch dimension and ensure mix is on the correct device

            est_s, _ = model(mix[None,], sources, do_test=True, pretrain=True)
            forward_end_time = datetime.now()  # Record forward pass end time
            forward_elapsed_time = (forward_end_time - forward_start_time).total_seconds()
            logger.debug(f"Batch {batch_idx} forward pass completed in {forward_elapsed_time:.2f} seconds.")

            sources = sources.squeeze(0)[None, :, 0, :]  # Remove batch dimension and use a single channel

            est_s = est_s[:, :, 0, :]  # Ensure the estimated sources have the correct dimensions
            reordered_s = loss_func(est_s, sources, return_est=True)[-1]

            mix_np = mix[[0], :].cpu().data.numpy()
            sources_np = sources.squeeze(0).cpu().data.numpy()
            reordered_s_np = reordered_s.squeeze(0).cpu().data.numpy()

            metrics_start_time = datetime.now()  # Record metrics computation start time
            sgs_metrics = get_metrics(mix_np, sources_np, reordered_s_np,
                                      sample_rate=8000,
                                      metrics_list=compute_metrics)
            metrics_end_time = datetime.now()  # Record metrics computation end time
            metrics_elapsed_time = (metrics_end_time - metrics_start_time).total_seconds()
            logger.debug(f"Batch {batch_idx} metrics computation completed in {metrics_elapsed_time:.2f} seconds.")

            all_metrics.append(pd.Series(sgs_metrics))

            end_time = datetime.now()  # Record end time
            elapsed_time = (end_time - start_time).total_seconds()
            logger.debug(f"Batch {batch_idx} processed in {elapsed_time:.2f} seconds.")

    avg_metrics = pd.DataFrame(all_metrics).mean()
    avg_sdr = avg_metrics['sdr']
    logger.info(f"Average SDR: {avg_sdr}")

    return avg_sdr


def export_fp32_onnx_model(model, example_inputs, save_path, opset_version=14):
    """Export the FP32 PyTorch model to an ONNX model."""
    logger.info(f"Exporting model to ONNX with inputs of shape: {[x.shape for x in example_inputs]}")
    torch_to_fp32_onnx(
        pt_fp32_model=model,
        save_path=save_path,
        example_inputs=example_inputs,
        opset_version=opset_version,
        verbose=True,
    )
    logger.info(f"Model exported to {save_path} successfully.")

def quantize_onnx_model(onnx_model_path, quantized_model_path, q_config, calib_dataloader):
    """Quantize the ONNX model using Intel Neural Compressor."""
    from neural_compressor.quantization import fit
    from neural_compressor.config import PostTrainingQuantConfig
    from neural_compressor.data import DataLoader

    tuning_criterion = TuningCriterion(max_trials=100)
    accuracy_criterion = AccuracyCriterion(higher_is_better=False, tolerable_loss=0.01)
    quantization_config = PostTrainingQuantConfig(
        approach="static",
        tuning_criterion=tuning_criterion,
        accuracy_criterion=accuracy_criterion,
        device="cpu",
        backend="onnxrt_qdq",
        op_type_dict = {
            "aten::linalg_solve": FP32,  # 尝试使用全称
            "linalg_solve": FP32,
            "linalg": FP32,  # 更泛化地覆盖 linalg 操作
            "Linear": FP32,
            "GroupNorm": FP32,
            "LayerNorm": FP32,
    }
    )

    q_model = fit(
        model=onnx_model_path,
        conf=quantization_config,
        calib_dataloader=calib_dataloader,
        eval_func=eval_func,
    )

    q_model.save(quantized_model_path)


def main(conf):  # 将 current_time 作为参数传递给 main 函数
    global model, val_loader1, model_device

    # 确保所有操作都在 CPU 上
    model_device = torch.device('cpu')

    # Define model and optimizer
    print("Creating model and optimizer")
    model, optimizer = make_model_and_optimizer(conf)

    # model = model.cuda()
    model_device = torch.device('cpu')  # 使用 CPU 作为设备
    # model_device = next(model.parameters()).device

    # Logging device info
    # logger.info(f"CUDA Available: {torch.cuda.is_available()}")
    logger.info(f"Using device: {model_device}")

    # model.to(model_device)  # 移动模型到适当的设备
    exp_dir = conf['main_args']['exp_dir']
    best_ckpt = conf['main_args']['best_ckpt']

    # 加载最好的模型
    if best_ckpt:
        checkpoint = torch.load(best_ckpt, map_location=model_device)  # 在CPU上加载模型
        if 'state_dict' in checkpoint:
            model.load_state_dict(checkpoint['state_dict'], strict=False)
            print(f"Loaded model from {best_ckpt}")
        else:
            print("Checkpoint does not contain 'state_dict'.")
    else:
        print("No best checkpoint provided.")

    train_loader, val_loader = make_dataloaders(**conf['data'], **conf['training'], channels=slice(0, 4))

    # 截取数据集的前, 100个
    def truncate_loader(loader, percentage=0.1):
        dataset = loader.dataset
        original_len = len(dataset)
        truncated_len =  int(original_len * percentage)
        indices = list(range(truncated_len))
        return torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False,
                                           num_workers=20,
                                           sampler=torch.utils.data.SubsetRandomSampler(indices))

    val_loader1 = truncate_loader(val_loader, percentage=0.03)

    os.makedirs(exp_dir, exist_ok=True)
    conf_path = os.path.join(exp_dir, 'conf.yml')
    with open(conf_path, 'w') as outfile:
        yaml.safe_dump(conf, outfile)

    logger.info("Starting the ONNX export and quantization process...")

    example_x = torch.randn(1, 4, 32000).to(model_device)  # 确保输入在 CPU 上
    example_s = torch.randn(1, 2, 4, 32000).to(model_device)  # 确保输入在 CPU 上

    try:

        # 调试模型输出
        outputs = model(example_x, example_s, do_test=False, pretrain=True)
        # print(f"Model outputs: {outputs}")

        # Step 1: Export the FP32 PyTorch model to an ONNX model
        onnx_model_path = "./saved_model/fp32_model.onnx"
        # 示例输入，假设batch_size=1, n_channels=4, n_samples=32000

        export_fp32_onnx_model(model, (example_x, example_s), save_path=onnx_model_path)
        logger.info(f"FP32 ONNX model has been successfully exported to {onnx_model_path}")

        # Step 2: Quantize the exported ONNX model
        quantized_onnx_model_path = "./saved_model/quantized_model.onnx"
        quantize_onnx_model(onnx_model_path, quantized_onnx_model_path, conf, val_loader1)
        logger.info(f"Quantized ONNX model has been successfully saved to {quantized_onnx_model_path}")
    except Exception as e:
        logger.error(f"Error during the ONNX export and quantization process: {str(e)}")

    logger.info("ONNX export and quantization process completed successfully.")

```

Given that I'm not directly using linalg_solve, I'm puzzled by why this error is occurring. Could it be related to the use of torch.linalg.matrix_norm or some other part of the code being translated during the ONNX export? 
Any guidance or suggestions on how to resolve or work around this issue would be greatly appreciated.


### Versions

PyTorch version: 2.2.0+cpu                   
Is debug build: False                        
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A              

OS: Microsoft Windows 10 教育版
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.8.0 (default, Nov  6 2019, 16:00:02) [MSC v.1916 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19041-SP0
Is CUDA available: False
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4090
Nvidia driver version: 527.56
cuDNN version: F:\CUDA\v11.8\bin\cudnn_ops_train64_8.dll
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=3500
DeviceID=CPU0
Family=205
L2CacheSize=8192
L2CacheSpeed=
Manufacturer=GenuineIntel
MaxClockSpeed=3500
Name=13th Gen Intel(R) Core(TM) i5-13600KF
ProcessorType=3
Revision=

Versions of relevant libraries:
[pip3] numpy==1.24.4
[pip3] onnx==1.16.0
[pip3] onnxruntime==1.16.0
[pip3] onnxscript==0.1.0.dev20240808
[pip3] pytorch-lightning==2.3.0
[pip3] pytorch-ranger==0.1.1
[pip3] torch==2.2.0+cpu
[pip3] torch-complex==0.4.3
[pip3] torch-optimizer==0.1.0
[pip3] torch-stoi==0.2.1
[pip3] torch-tb-profiler==0.4.3
[pip3] torchaudio==2.2.0+cpu
[pip3] torchmetrics==0.11.4
[pip3] torchvision==0.17.0+cpu
[conda] cudatoolkit               11.8.0               hd77b12b_0    defaults
[conda] libblas                   3.9.0                     8_mkl    conda-forge
[conda] libcblas                  3.9.0                     8_mkl    conda-forge
[conda] liblapack                 3.9.0                     8_mkl    conda-forge
[conda] mkl                       2021.4.0                 pypi_0    pypi
[conda] numpy                     1.24.4                   pypi_0    pypi
[conda] pytorch-cuda              11.8                 h24eeafa_5    pytorch
[conda] pytorch-lightning         2.3.0                    pypi_0    pypi
[conda] pytorch-ranger            0.1.1                    pypi_0    pypi
[conda] torch                     2.2.0+cpu                pypi_0    pypi
[conda] torch-complex             0.4.3                    pypi_0    pypi
[conda] torch-optimizer           0.1.0                    pypi_0    pypi
[conda] torch-stoi                0.2.1                    pypi_0    pypi
[conda] torch-tb-profiler         0.4.3                    pypi_0    pypi
[conda] torchaudio                2.2.0+cpu                pypi_0    pypi
[conda] torchmetrics              0.11.4                   pypi_0    pypi
[conda] torchvision               0.17.0+cpu               pypi_0    pypi