Firedrake

pyop3/__init__.py

@@ -3,14 +3,22 @@
 # tracebacks for @property methods so we remove it here.
 import pytools
 
-del pytools.RecordWithoutPickling.__getattr__
+try:
+    del pytools.RecordWithoutPickling.__getattr__
+except AttributeError:
+    pass
 del pytools
 
 
-import pyop3.transforms
-from pyop3.array import Array, HierarchicalArray, MultiArray, PetscMat
-from pyop3.axtree import Axis, AxisComponent, AxisTree  # noqa: F401
-from pyop3.buffer import DistributedBuffer  # noqa: F401
+import pyop3.ir
+import pyop3.transform
+from pyop3.array import Array, HierarchicalArray, MultiArray
+
+# TODO where should these live?
+from pyop3.array.harray import AxisVariable
+from pyop3.array.petsc import PetscMat, PetscMatAIJ, PetscMatPreallocator  # noqa: F401
+from pyop3.axtree import Axis, AxisComponent, AxisTree, PartialAxisTree  # noqa: F401
+from pyop3.buffer import DistributedBuffer, NullBuffer  # noqa: F401
 from pyop3.dtypes import IntType, ScalarType  # noqa: F401
 from pyop3.itree import (  # noqa: F401
     AffineSliceComponent,
@@ -23,18 +31,25 @@
     Subset,
     TabulatedMapComponent,
 )
+from pyop3.itree.tree import ScalarIndex
 from pyop3.lang import (  # noqa: F401
     INC,
     MAX_RW,
     MAX_WRITE,
     MIN_RW,
     MIN_WRITE,
+    NA,
     READ,
     RW,
     WRITE,
+    AddAssignment,
+    DummyKernelArgument,
     Function,
     Loop,
+    OpaqueKernelArgument,
+    Pack,
+    ReplaceAssignment,
     do_loop,
     loop,
 )
-from pyop3.tensor import Dat, Global, Mat, Tensor  # noqa: F401
+from pyop3.sf import StarForest, serial_forest, single_star

pyop3/array/__init__.py

@@ -1,7 +1,10 @@
+# arguably put this directly in pyop3/__init__.py
+# no use namespacing here really
+
 from .base import Array  # noqa: F401
 from .harray import (  # noqa: F401
     ContextSensitiveMultiArray,
     HierarchicalArray,
     MultiArray,
 )
-from .petsc import PackedPetscMatAIJ, PetscMat, PetscMatAIJ  # noqa: F401
+from .petsc import PetscMat, PetscMatAIJ  # noqa: F401

pyop3/array/base.py

@@ -1,6 +1,6 @@
 import abc
 
-from pyop3.lang import KernelArgument
+from pyop3.lang import KernelArgument, ReplaceAssignment
 from pyop3.utils import UniqueNameGenerator
 
 
@@ -13,7 +13,5 @@ def __init__(self, name=None, *, prefix=None) -> None:
             raise ValueError("Can only specify one of name and prefix")
         self.name = name or self._name_generator(prefix or self._prefix)
 
-    @property
-    @abc.abstractmethod
-    def valid_ranks(self):
-        pass
+    def assign(self, other):
+        return ReplaceAssignment(self, other)