Merge branch 'master' into dask

dmentipl · Apr 14, 2020 · a8fa457 · a8fa457
2 parents ba193d7 + b9a49bc
commit a8fa457
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/plonk/snap/snap.py b/plonk/snap/snap.py
@@ -946,7 +946,10 @@ def _get_array(self, name: str, sinks: bool = False) -> ndarray:
         if name in self._array_registry or name in self._sink_registry:
             array = self._get_array_from_registry(name, sinks)
             if self.cache_arrays:
-                self._arrays[name] = array
+                if sinks:
+                    self._sinks[name] = array
+                else:
+                    self._arrays[name] = array
             if index is None:
                 return array
             return array[:, index]