UniPC and `linalg_solve_ex` with AuraFlow and ComfyUI

Hopefully the right place to be here…
I’ve been encountering an error while using Pytorch in conjunction with AuraFlow and UniPC which brought me to the support matrix. I noticed there that linalg_solve_triangular is marked as green and enabled, but I don’t see linalg_solve_ex listed there. The instructions
suggest checking the matrix first before commenting on the GitHub issue.

My questions are:

  1. Are solve_ex and solve_triangular different operations, or are they both part of the same linalg_solve functionality?
  2. Should I comment on the GitHub issue to request support for linalg_solve_ex, even though linalg_solve_triangular is already supported?

Exact Error

got prompt
  8%|██████████▉                                                                                                                              | 2/25 [00:15<02:52,  7.52s/it]
!!! Exception during processing !!! The operator 'aten::_linalg_solve_ex.result' is not currently implemented for the MPS device. If you want this op to be considered for addition please comment on https://github.com/pytorch/pytorch/issues/141287 and mention use-case, that resulted in missing op as well as commit hash 3dd667797a7f46f6c241e492ad1586058acc5174. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.
Traceback (most recent call last):
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/execution.py", line 324, in execute
    output_data, output_ui, has_subgraph = get_output_data(obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/execution.py", line 199, in get_output_data
    return_values = _map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/execution.py", line 170, in _map_node_over_list
    process_inputs(input_dict, i)
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/execution.py", line 159, in process_inputs
    results.append(getattr(obj, func)(**inputs))
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/nodes.py", line 1467, in sample
    return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/nodes.py", line 1434, in common_ksampler
    samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/sample.py", line 43, in sample
    samples = sampler.sample(noise, positive, negative, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/samplers.py", line 1020, in sample
    return sample(self.model, noise, positive, negative, cfg, self.device, sampler, sigmas, self.model_options, latent_image=latent_image, denoise_mask=denoise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/samplers.py", line 918, in sample
    return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/samplers.py", line 904, in sample
    output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/patcher_extension.py", line 110, in execute
    return self.original(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/samplers.py", line 873, in outer_sample
    output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/samplers.py", line 857, in inner_sample
    samples = executor.execute(self, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/patcher_extension.py", line 110, in execute
    return self.original(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/samplers.py", line 714, in sample
    samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/extra_samplers/uni_pc.py", line 870, in sample_unipc
    x = uni_pc.sample(noise, timesteps=timesteps, skip_type="time_uniform", method="multistep", order=order, lower_order_final=True, callback=callback, disable_pbar=disable)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/extra_samplers/uni_pc.py", line 724, in sample
    x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, init_order, use_corrector=True)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/extra_samplers/uni_pc.py", line 472, in multistep_uni_pc_update
    return self.multistep_uni_pc_bh_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/unauthorized/Documents/GitHub/src/comfyui/comfy/extra_samplers/uni_pc.py", line 654, in multistep_uni_pc_bh_update
    rhos_c = torch.linalg.solve(R, b)
             ^^^^^^^^^^^^^^^^^^^^^^^^

M4 Max 16 Nov 2024 Sequoia 128gb 15.1.1
Python 3.12.7 [Clang 18.1.8 ] on darwin
Using Python 3.12.7 environment .virtualenvs/comfyui
Package Version


aiohappyeyeballs 2.4.4
aiohttp 3.11.10
aiosignal 1.3.1
attrs 24.2.0
certifi 2024.8.30
cffi 1.17.1
charset-normalizer 3.4.0
einops 0.8.0
filelock 3.16.1
frozenlist 1.5.0
fsspec 2024.10.0
huggingface-hub 0.26.5
idna 3.10
jinja2 3.1.4
kornia 0.7.4
kornia-rs 0.1.7
markupsafe 2.1.5
mpmath 1.3.0
multidict 6.1.0
networkx 3.4.2
numpy 2.1.2
packaging 24.2
pillow 11.0.0
propcache 0.2.1
psutil 6.1.0
pycparser 2.22
pyyaml 6.0.2
regex 2024.11.6
requests 2.32.3
safetensors 0.4.5
scipy 1.14.1
sentencepiece 0.2.0
setuptools 70.2.0
soundfile 0.12.1
spandrel 0.4.0
sympy 1.13.1
tokenizers 0.21.0
torch 2.6.0.dev20241206
torchaudio 2.5.0.dev20241206
torchsde 0.2.6
torchvision 0.20.0.dev20241206
tqdm 4.67.1
trampoline 0.1.2
transformers 4.47.0
typing-extensions 4.12.2
urllib3 2.2.3
yarl 1.18.3

Any insights or suggestions would be greatly appreciated!

Thanks