remove MPI

Author: ferrine

pytensor/tensor/io.py

@@ -1,10 +1,9 @@
 import numpy as np
 
-from pytensor.graph.basic import Apply, Constant, Variable
+from pytensor.graph.basic import Apply, Constant
 from pytensor.graph.op import Op
 from pytensor.link.c.type import Generic
 from pytensor.tensor.type import tensor
-from pytensor.utils import key_to_cmp
 
 
 class LoadFromDisk(Op):
@@ -92,229 +91,4 @@ def load(path, dtype, shape, mmap_mode=None):
     return LoadFromDisk(dtype, shape, mmap_mode)(path)
 
 
-##########################
-# MPI
-##########################
-
-try:
-    from mpi4py import MPI
-except ImportError:
-    mpi_enabled = False
-else:
-    comm = MPI.COMM_WORLD
-    mpi_enabled = True
-
-
-class MPIRecv(Op):
-    """
-    An operation to asynchronously receive an array to a remote host using MPI.
-
-    See Also
-    --------
-    MPIRecv
-    MPIWait
-
-    Notes
-    -----
-    Non-differentiable.
-
-    """
-
-    __props__ = ("source", "tag", "shape", "dtype")
-
-    def __init__(self, source, tag, shape, dtype):
-        self.source = source
-        self.tag = tag
-        self.shape = shape
-        self.dtype = np.dtype(dtype)  # turn "float64" into numpy.float64
-        self.static_shape = (None,) * len(shape)
-
-    def make_node(self):
-        return Apply(
-            self,
-            [],
-            [
-                Variable(Generic(), None),
-                tensor(self.dtype, shape=self.static_shape),
-            ],
-        )
-
-    def perform(self, node, inp, out):
-
-        data = np.zeros(self.shape, dtype=self.dtype)
-        request = comm.Irecv(data, self.source, self.tag)
-
-        out[0][0] = request
-        out[1][0] = data
-
-    def __str__(self):
-        return f"MPIRecv{{source: {int(self.source)}, tag: {int(self.tag)}, shape: {self.shape}, dtype: {self.dtype}}}"
-
-    def infer_shape(self, fgraph, node, shapes):
-        return [None, self.shape]
-
-    def do_constant_folding(self, fgraph, node):
-        return False
-
-
-class MPIRecvWait(Op):
-    """
-    An operation to wait on a previously received array using MPI.
-
-    See Also
-    --------
-    MPIRecv
-
-    Notes
-    -----
-    Non-differentiable.
-
-    """
-
-    __props__ = ("tag",)
-
-    def __init__(self, tag):
-        self.tag = tag
-
-    def make_node(self, request, data):
-        return Apply(
-            self,
-            [request, data],
-            [tensor(data.dtype, shape=data.type.shape)],
-        )
-
-    def perform(self, node, inp, out):
-
-        request = inp[0]
-        data = inp[1]
-
-        request.wait()
-
-        out[0][0] = data
-
-    def infer_shape(self, fgraph, node, shapes):
-        return [shapes[1]]
-
-    view_map = {0: [1]}
-
-
-class MPISend(Op):
-    """
-    An operation to asynchronously Send an array to a remote host using MPI.
-
-    See Also
-    --------
-    MPIRecv
-    MPISendWait
-
-    Notes
-    -----
-    Non-differentiable.
-
-    """
-
-    __props__ = ("dest", "tag")
-
-    def __init__(self, dest, tag):
-        self.dest = dest
-        self.tag = tag
-
-    def make_node(self, data):
-        return Apply(self, [data], [Variable(Generic(), None), data.type()])
-
-    view_map = {1: [0]}
-
-    def perform(self, node, inp, out):
-
-        data = inp[0]
-
-        request = comm.Isend(data, self.dest, self.tag)
-
-        out[0][0] = request
-        out[1][0] = data
-
-    def __str__(self):
-        return f"MPISend{{dest: {int(self.dest)}, tag: {int(self.tag)}}}"
-
-
-class MPISendWait(Op):
-    """
-    An operation to wait on a previously sent array using MPI.
-
-    See Also
-    --------
-    MPISend
-
-    Notes
-    -----
-    Non-differentiable.
-
-    """
-
-    __props__ = ("tag",)
-
-    def __init__(self, tag):
-        self.tag = tag
-
-    def make_node(self, request, data):
-        return Apply(self, [request, data], [Variable(Generic(), None)])
-
-    def perform(self, node, inp, out):
-        request = inp[0]
-        request.wait()
-        out[0][0] = True
-
-
-def isend(var, dest, tag):
-    """
-    Non blocking send.
-    """
-    return MPISend(dest, tag)(var)
-
-
-def send(var, dest, tag):
-    """
-    Blocking send.
-    """
-    return MPISendWait(tag)(*isend(var, dest, tag))
-
-
-def irecv(shape, dtype, source, tag):
-    """
-    Non-blocking receive.
-    """
-    return MPIRecv(source, tag, shape, dtype)()
-
-
-def recv(shape, dtype, source, tag):
-    """
-    Blocking receive.
-    """
-    return MPIRecvWait(tag)(*irecv(shape, dtype, source, tag))
-
-
-# Ordering keys for scheduling
-def mpi_send_wait_key(a):
-    """Wait as long as possible on Waits, Start Send/Recvs early."""
-    if isinstance(a.op, (MPIRecvWait, MPISendWait)):
-        return 1
-    if isinstance(a.op, (MPIRecv, MPISend)):
-        return -1
-    return 0
-
-
-def mpi_tag_key(a):
-    """Break MPI ties by using the variable tag - prefer lower tags first."""
-    if isinstance(a.op, (MPISend, MPIRecv, MPIRecvWait, MPISendWait)):
-        return a.op.tag
-    else:
-        return 0
-
-
-mpi_send_wait_cmp = key_to_cmp(mpi_send_wait_key)
-mpi_tag_cmp = key_to_cmp(mpi_tag_key)
-
-mpi_keys = (mpi_send_wait_key, mpi_tag_key)
-mpi_cmps = (mpi_send_wait_cmp, mpi_tag_cmp)
-
 __all__ = ["load"]

tests/link/test_link.py

@@ -4,15 +4,13 @@
 import numpy as np
 
 import pytensor
-from pytensor.compile.mode import Mode
 from pytensor.graph import fg
 from pytensor.graph.basic import Apply, Constant, Variable, clone
 from pytensor.graph.op import Op
 from pytensor.graph.type import Type
 from pytensor.link.basic import Container, Linker, PerformLinker, WrapLinker
-from pytensor.link.c.basic import OpWiseCLinker
-from pytensor.tensor.type import matrix, scalar
-from pytensor.utils import cmp, to_return_values
+from pytensor.tensor.type import scalar
+from pytensor.utils import to_return_values
 
 
 def make_function(linker: Linker, unpack_single: bool = True, **kwargs) -> Callable:
@@ -219,26 +217,6 @@ def wrap(fgraph, i, node, th):
         assert o[0].data == 1.5
 
 
-def test_sort_schedule_fn():
-    from pytensor.graph.sched import make_depends, sort_schedule_fn
-
-    x = matrix("x")
-    y = pytensor.tensor.dot(x[:5] * 2, x.T + 1).T
-
-    def str_cmp(a, b):
-        return cmp(str(a), str(b))  # lexicographical sort
-
-    linker = OpWiseCLinker(schedule=sort_schedule_fn(str_cmp))
-    mode = Mode(linker=linker)
-    f = pytensor.function((x,), (y,), mode=mode)
-
-    nodes = f.maker.linker.make_all()[-1]
-    depends = make_depends()
-    for a, b in zip(nodes[:-1], nodes[1:]):
-        if not depends((b, a)):
-            assert str(a) < str(b)
-
-
 def test_container_deepcopy():
     # This is a test to a work around a NumPy bug.