diff --git a/src/gt4py/next/common.py b/src/gt4py/next/common.py index 9b2870e1c0..4fc155951e 100644 --- a/src/gt4py/next/common.py +++ b/src/gt4py/next/common.py @@ -574,6 +574,12 @@ def replace(self, index: int | Dimension, *named_ranges: NamedRange) -> Domain: return Domain(dims=dims, ranges=ranges) + def __getstate__(self): + state = self.__dict__.copy() + # remove cached property + state.pop("slice_at", None) + return state + FiniteDomain: TypeAlias = Domain[FiniteUnitRange] diff --git a/tests/next_tests/regression_tests/embedded_tests/test_domain_pickle.py b/tests/next_tests/regression_tests/embedded_tests/test_domain_pickle.py new file mode 100644 index 0000000000..1eb12ec98c --- /dev/null +++ b/tests/next_tests/regression_tests/embedded_tests/test_domain_pickle.py @@ -0,0 +1,13 @@ +import pickle + +from gt4py.next import common + +I = common.Dimension("I") +J = common.Dimension("J") + +def test_domain_pickle_after_slice(): + domain = common.domain(((I, (2, 4)), (J, (3, 5)))) + # use slice_at to populate cached property + domain.slice_at[2:5, 5:7] + + pickle.dumps(domain) \ No newline at end of file