Issues with Cellpose on macOS using the MPS Backend

[Vijayishwerj]

Hello Cellpose Team,

I am experiencing several issues while using Cellpose for segmentation and training models on my Mac Studio M2 Ultra. I would appreciate your guidance to resolve these problems. Below are the details:

Environment Information

•	Operating System: macOS
•	Device: Mac Studio M2 Ultra
•	Python Version: 3.11.10
•	Cellpose Version: 3.1.0
•	Torch Version: 2.5.1
•	Backend: MPS (Metal Performance Shaders)

Issues Faced

1.	GPU Incompatibility with Sparse Tensor Operations:
•	Error: NotImplementedError: Could not run 'aten::_sparse_coo_tensor_with_dims_and_tensors' with arguments from the 'SparseMPS' backend....
•	This suggests that the MPS backend lacks support for sparse tensor operations, leading to failures during segmentation.
2.	Fallback to CPU:
•	When the MPS backend fails, computation falls back to the CPU.
•	However, warnings about the missing MKL optimizations slow down performance significantly:  WARNING: MKL version on torch not working/installed - CPU version will be slightly slower.


3.	GPU Training Issues:
•	The latest version of Cellpose mandates GPU use for training. However, the MPS backend does not complete tasks, and training without a GPU seems impossible.
4.	Performance Bottleneck:
•	The fallback to CPU leads to extremely long training times, making it impractical for large datasets.

Run Logs

Attached below is the terminal output with verbose mode enabled, showing the errors and relevant information:
• MPS backend available: torch.backends.mps.is_available() returns True.
• The operation fails during sparse tensor computation (aten::_sparse_coo_tensor_with_dims_and_tensors).

NotImplementedError: Could not run 'aten::_sparse_coo_tensor_with_dims_and_tensors' with arguments from the 'SparseMPS' backend. ####

1.	Verified that the MPS backend is correctly installed and available.
2.	Updated to the latest versions of Cellpose, Python, and PyTorch.
3.	Attempted fallback to the CPU but encountered performance issues due to lack of MKL support.
4.	Explored replacing sparse operations with dense ones but encountered compatibility constraints.

Request for Assistance

1.	GPU Support:
•	Are there plans to improve sparse tensor compatibility for the MPS backend in future versions?