Don't compress PyTorch matrices when requires_grad

This is due to PyTorch not supporting grad of bincount
tum-pbs · Jan 15, 2025 · 5bf9d58 · 5bf9d58
1 parent 14274b5
commit 5bf9d58
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/phiml/math/_trace.py b/phiml/math/_trace.py
@@ -728,6 +728,8 @@ def matrix_from_function(f: Callable,
     else:
         matrix, bias = tracer_to_coo(tracer, sparsify, separate_independent)
     # --- Compress ---
+    if matrix.backend.name == 'torch' and matrix._values._native.requires_grad:
+        auto_compress = False  # PyTorch doesn't support gradient of bincount (used in compression)
     if auto_compress and matrix.backend.supports(Backend.mul_csr_dense) and target_backend.supports(Backend.mul_csr_dense) and isinstance(matrix, SparseCoordinateTensor):
         matrix = matrix.compress_rows()
     # elif backend.supports(Backend.mul_csc_dense):