Merge pull request #2435 from devitocodes/sparse-pickle

compiler: fix unevaluation of sparse cluster
devitocodes · Aug 1, 2024 · 4e95e2f · 4e95e2f
2 parents 97c3526 + 03eb51b
commit 4e95e2f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/devito/passes/clusters/unevaluate.py b/devito/passes/clusters/unevaluate.py
@@ -7,7 +7,7 @@
 __all__ = ['unevaluate']
 
 
-@cluster_pass
+@cluster_pass(mode='all')
 def unevaluate(cluster):
     exprs = [_unevaluate(e) for e in cluster.exprs]
 

diff --git a/tests/test_pickle.py b/tests/test_pickle.py
@@ -864,7 +864,8 @@ def test_full_model(self, pickle):
         assert np.isclose(np.linalg.norm(ricker.data), np.linalg.norm(new_ricker.data))
         # FIXME: fails randomly when using data.flatten() AND numpy is using MKL
 
-    def test_usave_sampled(self, pickle):
+    @pytest.mark.parametrize('subs', [False, True])
+    def test_usave_sampled(self, pickle, subs):
         grid = Grid(shape=(10, 10, 10))
         u = TimeFunction(name="u", grid=grid, time_order=2, space_order=8)
 
@@ -884,7 +885,9 @@ def test_usave_sampled(self, pickle):
 
         eqn = [stencil] + src_term
         eqn += [Eq(u0_save, u)]
-        op_fwd = Operator(eqn)
+
+        subs = grid.spacing_map if subs else {}
+        op_fwd = Operator(eqn, subs=subs)
 
         tmp_pickle_op_fn = "tmp_operator.pickle"
         pickle.dump(op_fwd, open(tmp_pickle_op_fn, "wb"))