Rename project and switch to PyTensor

Author: michaelosthege

.github/workflows/test.yml

@@ -45,10 +45,10 @@ jobs:
         python-version: ${{matrix.python-version}}
         use-mamba: true
         use-only-tar-bz2: true # IMPORTANT: This needs to be set for caching to work properly!
-    - name: Install aesara-federated
+    - name: Install pytensor-federated
       run: |
           conda activate aefenv
           pip install -e .
     - name: Run tests
       run: |
-        pytest -v --cov=./aesara_federated --cov-report xml --cov-report term-missing .
+        pytest -v --cov=./pytensor_federated --cov-report xml --cov-report term-missing .

.pre-commit-config.yaml

@@ -15,9 +15,9 @@ repos:
     - id: isort
       name: isort
       args: ["--profile", "black"]
-      exclude: (aesara_federated/rpc.py|aesara_federated/npproto/)
+      exclude: (pytensor_federated/rpc.py|pytensor_federated/npproto/)
 - repo: https://github.com/psf/black
   rev: 22.3.0
   hooks:
     - id: black
-      exclude: (aesara_federated/rpc.py|aesara_federated/npproto/)
+      exclude: (pytensor_federated/rpc.py|pytensor_federated/npproto/)

README.md

@@ -1,12 +1,12 @@
-[![pipeline](https://github.com/michaelosthege/aesara-federated/workflows/test/badge.svg)](https://github.com/michaelosthege/aesara-federated/actions)
+[![pipeline](https://github.com/michaelosthege/pytensor-federated/workflows/test/badge.svg)](https://github.com/michaelosthege/pytensor-federated/actions)
 
-# `aesara-federated`
-This package implements federated computing with [Aesara](https://github.com/aesara-devs/aesara).
+# `pytensor-federated`
+This package implements federated computing with [PyTensor](https://github.com/pymc-devs/pytensor).
 
-Using `aesara-federated`, differentiable cost functions can be computed on federated nodes.
+Using `pytensor-federated`, differentiable cost functions can be computed on federated nodes.
 Inputs and outputs are transmitted in binary via a bidirectional gRPC stream.
 
-A client side `LogpGradOp` is provided to conveniently embed federated compute operations in Aesara graphs such as a [PyMC](https://github.com/pymc-devs/pymc) model.
+A client side `LogpGradOp` is provided to conveniently embed federated compute operations in PyTensor graphs such as a [PyMC](https://github.com/pymc-devs/pymc) model.
 
 The example code implements a simple Bayesian linear regression to data that is "private" to the federated compute process.
 
@@ -21,7 +21,7 @@ python demo_model.py
 ```
 
 ## Architecture
-`aesara-federated` is designed to be a very generalizable framework for federated computing with gRPC, but it comes with implementations for Aesara, and specifically for use cases of Bayesian inference.
+`pytensor-federated` is designed to be a very generalizable framework for federated computing with gRPC, but it comes with implementations for PyTensor, and specifically for use cases of Bayesian inference.
 This is reflected in the actual implementation, where the most basic gRPC service implementation -- the `ArraysToArraysService` -- is wrapped by a few implementation flavors, specifically for common use cases in Bayesian inference.
 
 At the core, everything is built around an `ArraysToArrays` gRPC service, which takes any number of (NumPy) arrays as parameters, and returns any number of (NumPy) arrays as outputs.
@@ -54,10 +54,10 @@ Different sub-graphs of this example could be wrapped by an `ArraysToArraysServi
 If the entire model is differentiable, one can even return gradients.
 For example, with a linear model: `[slope, intercept] -> [LL, dLL_dslope, dLL_dintercept]`.
 
-The role of Aesara here is purely technical:
-Aesara is a graph computation framework that implements auto-differentiation.
-Wrapping the `ArraysToArraysServiceClient` in Aesara `Op`s simply makes it easier to build more sophisticated compute graphs.
-Aesara is also the computatation backend for PyMC, which is the most popular framework for Bayesian inference in Python.
+The role of PyTensor here is purely technical:
+PyTensor is a graph computation framework that implements auto-differentiation.
+Wrapping the `ArraysToArraysServiceClient` in PyTensor `Op`s simply makes it easier to build more sophisticated compute graphs.
+PyTensor is also the computatation backend for PyMC, which is the most popular framework for Bayesian inference in Python.
 
 
 ## Installation & Contributing

demo_model.py

@@ -6,7 +6,7 @@
 import numpy as np
 import pymc as pm
 
-from aesara_federated import AsyncLogpGradOp, LogpGradOp, LogpGradServiceClient
+from pytensor_federated import AsyncLogpGradOp, LogpGradOp, LogpGradServiceClient
 
 _log = logging.getLogger(__file__)
 logging.basicConfig(level=logging.INFO)

demo_node.py

@@ -5,12 +5,12 @@
 import time
 from typing import Sequence, Tuple
 
-import aesara
-import aesara.tensor as at
 import grpclib
 import numpy as np
+import pytensor
+import pytensor.tensor as at
 
-from aesara_federated import ArraysToArraysService, wrap_logp_grad_func
+from pytensor_federated import ArraysToArraysService, wrap_logp_grad_func
 
 _log = logging.getLogger(__file__)
 logging.basicConfig(level=logging.INFO)
@@ -36,7 +36,7 @@ def _make_function(x, y, sigma):
         pdf = 1 / (sigma * np.sqrt(2 * np.pi)) * at.exp(-0.5 * ((y - pred) / sigma) ** 2)
         logp = at.log(pdf).sum()
         grad = at.grad(logp, wrt=[intercept, slope])
-        fn = aesara.function(
+        fn = pytensor.function(
             inputs=[intercept, slope],
             outputs=[logp, *grad],
         )

environment.yml

@@ -9,4 +9,4 @@ dependencies:
   - pip
   - pip:
     - betterproto[compiler]==2.0.0b5
-    - pymc>=4.1.6
+    - pymc==5.0.0

protobufs/generate.py

@@ -8,7 +8,7 @@
 
 DP_HERE = pathlib.Path(__file__).parent.absolute()
 DP_ROOT = DP_HERE.parent
-DP_PROJ = DP_ROOT / "aesara_federated"
+DP_PROJ = DP_ROOT / "pytensor_federated"
 FP_PROTO = DP_HERE / "service.proto"
 
 

pyproject.toml

@@ -6,6 +6,6 @@ line-length = 100
 
 [tool.mypy]
 exclude = [
-  "^aesara_federated/test_*",
+  "^pytensor_federated/test_*",
 ]
 ignore_missing_imports = true

aesara_federated/__init__.py renamed to pytensor_federated/__init__.py

@@ -19,4 +19,4 @@
 from .service import ArraysToArraysService, ArraysToArraysServiceClient
 from .signatures import ComputeFunc, LogpFunc, LogpGradFunc
 
-__version__ = "0.3.0"
+__version__ = "0.4.0"

aesara_federated/common.py renamed to pytensor_federated/common.py

@@ -1,14 +1,14 @@
 import asyncio
 from typing import Any, Callable, List, Optional, Sequence
 
-import aesara.tensor as at
-from aesara.compile import optdb
-from aesara.compile.ops import FromFunctionOp
-from aesara.graph import FunctionGraph
-from aesara.graph.basic import Apply, Variable, is_in_ancestors
-from aesara.graph.features import ReplaceValidate
-from aesara.graph.op import Op, OutputStorageType, ParamsInputType
-from aesara.graph.rewriting.basic import GraphRewriter
+import pytensor.tensor as at
+from pytensor.compile import optdb
+from pytensor.compile.ops import FromFunctionOp
+from pytensor.graph import FunctionGraph
+from pytensor.graph.basic import Apply, Variable, is_in_ancestors
+from pytensor.graph.features import ReplaceValidate
+from pytensor.graph.op import Op, OutputStorageType, ParamsInputType
+from pytensor.graph.rewriting.basic import GraphRewriter
 
 from .utils import get_useful_event_loop
 
@@ -37,7 +37,7 @@ async def perform_async(
 
 
 class AsyncFromFunctionOp(AsyncOp, FromFunctionOp):
-    """Async version of the ``aesara.compile.ops.FromFunctionOp``.
+    """Async version of the ``pytensor.compile.ops.FromFunctionOp``.
 
     Note that ``AsyncOp.perform`` overrides ``FromFunctionOp.perform`` by MRO.
     """

aesara_federated/npproto/__init__.py renamed to pytensor_federated/npproto/__init__.py

@@ -3,8 +3,8 @@
 import numpy
 import pytest
 
-from aesara_federated import npproto
-from aesara_federated.npproto import utils
+from pytensor_federated import npproto
+from pytensor_federated.npproto import utils
 
 
 class TestUtils:

aesara_federated/npproto/utils.py renamed to pytensor_federated/npproto/utils.py

@@ -2,15 +2,15 @@
 import time
 from typing import Sequence, Tuple
 
-import aesara
-import aesara.tensor as at
 import numpy
+import pytensor
+import pytensor.tensor as at
 import pytest
-from aesara.graph import FunctionGraph
-from aesara.graph.basic import Apply, Variable
-from aesara.graph.op import Op
+from pytensor.graph import FunctionGraph
+from pytensor.graph.basic import Apply, Variable
+from pytensor.graph.op import Op
 
-from aesara_federated import op_async
+from pytensor_federated import op_async
 
 
 class _AsyncDelay(op_async.AsyncOp):
@@ -41,7 +41,7 @@ def test_perform(self):
         d = at.scalar()
         out = delay_op(d)
         # Compile a function to exclude compile time from delay measurement
-        f = aesara.function([d], [out])
+        f = pytensor.function([d], [out])
         ts = time.perf_counter()
         f(0.5)
         assert 0.5 < time.perf_counter() - ts < 0.6
@@ -60,12 +60,12 @@ async def _fn(x):
             otypes=[at.dscalar],
         )
         assert isinstance(affo, op_async.AsyncOp)
-        assert isinstance(affo, aesara.compile.ops.FromFunctionOp)
+        assert isinstance(affo, pytensor.compile.ops.FromFunctionOp)
 
         d = at.scalar()
         out = affo(d)
         # Compile a function to exclude compile time from delay measurement
-        f = aesara.function([d], [out])
+        f = pytensor.function([d], [out])
         ts = time.perf_counter()
         f(0.5)
         assert 0.5 < time.perf_counter() - ts < 0.6
@@ -97,7 +97,7 @@ def test_perform(self):
 
         # Evaluating the delays in parallel is faster than the sum of delays.
         # We do this with a compiled function to exclude compile time from delay measurement.
-        f = aesara.function([d1, d2], [dsum])
+        f = pytensor.function([d1, d2], [dsum])
         t_start = time.perf_counter()
         delay_sum = f(0.5, 0.2)[0]
         t_took = time.perf_counter() - t_start
@@ -152,10 +152,10 @@ def test_parallelize_async_applies():
 
 def _measure_fg(fg: FunctionGraph, *inputs) -> Tuple[Sequence[numpy.ndarray], float]:
     """Measure the runtime of a function compiled from `fg`."""
-    f = aesara.function(
+    f = pytensor.function(
         fg.inputs,
         fg.outputs,
-        mode=aesara.compile.FAST_COMPILE,  # skip the async fusion
+        mode=pytensor.compile.FAST_COMPILE,  # skip the async fusion
     )
     t0 = time.perf_counter()
     outputs = f(*inputs)
@@ -199,7 +199,7 @@ def test_fuse_asyncs_by_default():
     delay = _AsyncDelay()
     a, b = at.scalars("ab")
     c = delay(a) + delay(b)
-    f = aesara.function([a, b], [c])
+    f = pytensor.function([a, b], [c])
     t0 = time.perf_counter()
     f(0.25, 0.25)
     assert time.perf_counter() - t0 < 0.3

aesara_federated/op_async.py renamed to pytensor_federated/op_async.py

@@ -8,10 +8,10 @@
 import numpy as np
 import pytest
 
-from aesara_federated import service, signatures
-from aesara_federated.npproto.utils import ndarray_from_numpy, ndarray_to_numpy
-from aesara_federated.rpc import GetLoadResult, InputArrays, OutputArrays
-from aesara_federated.utils import get_useful_event_loop
+from pytensor_federated import service, signatures
+from pytensor_federated.npproto.utils import ndarray_from_numpy, ndarray_to_numpy
+from pytensor_federated.rpc import GetLoadResult, InputArrays, OutputArrays
+from pytensor_federated.utils import get_useful_event_loop
 
 
 def test_compute_function():

aesara_federated/rpc.py renamed to pytensor_federated/rpc.py

@@ -1,7 +1,7 @@
 import asyncio
 
-from aesara_federated import utils
-from aesara_federated.rpc import GetLoadResult
+from pytensor_federated import utils
+from pytensor_federated.rpc import GetLoadResult
 
 
 def test_argmin_load():

aesara_federated/service.py renamed to pytensor_federated/service.py

@@ -4,19 +4,19 @@
 import sys
 import time
 
-import aesara
-import aesara.tensor as at
 import arviz
 import grpclib
 import numpy as np
 import pymc as pm
+import pytensor
+import pytensor.tensor as at
 import pytest
 import scipy
-from aesara.compile.ops import FromFunctionOp
-from aesara.graph.basic import Apply, Variable
+from pytensor.compile.ops import FromFunctionOp
+from pytensor.graph.basic import Apply, Variable
 
-from aesara_federated import common, op_async, service, wrapper_ops
-from aesara_federated.utils import get_useful_event_loop
+from pytensor_federated import common, op_async, service, wrapper_ops
+from pytensor_federated.utils import get_useful_event_loop
 
 
 class _MockLogpGradOpClient:
@@ -228,7 +228,7 @@ def test_grad(self):
         b = at.scalar()
         logp, *_ = flop(a, b)
         ga, gb = at.grad(logp, [a, b])
-        fn = aesara.function(inputs=[a, b], outputs=[logp, ga, gb])
+        fn = pytensor.function(inputs=[a, b], outputs=[logp, ga, gb])
         exlogp, (exda, exdb) = dummy_quadratic_model(1.4, 0.5)
         actual = fn(1.4, 0.5)
         np.testing.assert_array_equal(actual[0], exlogp)

aesara_federated/signatures.py renamed to pytensor_federated/signatures.py

@@ -1,18 +1,18 @@
 from typing import Any, Callable, List, Optional, Sequence, Union
 
-import aesara
-import aesara.tensor as at
 import numpy as np
-from aesara.compile.ops import FromFunctionOp
-from aesara.graph.basic import Apply, Variable
-from aesara.graph.op import Op, OutputStorageType, ParamsInputType
+import pytensor
+import pytensor.tensor as at
+from pytensor.compile.ops import FromFunctionOp
+from pytensor.graph.basic import Apply, Variable
+from pytensor.graph.op import Op, OutputStorageType, ParamsInputType
 
 from .op_async import AsyncFromFunctionOp, AsyncOp
 from .signatures import ComputeFunc, LogpFunc, LogpGradFunc
 
 
 class ArraysToArraysOp(FromFunctionOp):
-    """Alias for the `aesara.compile.ops.FromFunctionOp`.
+    """Alias for the `pytensor.compile.ops.FromFunctionOp`.
 
     This alias exists for more convenient imports,
     more informative type hints,
@@ -124,7 +124,7 @@ def grad(self, inputs: Sequence[Variable], output_grads: List[Variable]) -> List
         # one w.r.t. logp
         g_logp, *gs_inputs = output_grads
         for i, g in enumerate(gs_inputs):
-            if not isinstance(g.type, aesara.gradient.DisconnectedType):
+            if not isinstance(g.type, pytensor.gradient.DisconnectedType):
                 raise ValueError(f"Can't propagate gradients wrt parameter {i+1}")
         # Call again on the original inputs, to obtain a handle
         # on the gradient. The computation will not actually be

aesara_federated/test_npproto.py renamed to pytensor_federated/test_npproto.py

@@ -4,7 +4,7 @@
 
 import setuptools
 
-__packagename__ = "aesara_federated"
+__packagename__ = "pytensor_federated"
 ROOT = pathlib.Path(__file__).parent
 
 
@@ -38,8 +38,8 @@ def get_version():
     packages=setuptools.find_packages(),
     version=__version__,
     description="This package helps to reduce the amount of boilerplate code when creating Airflow DAGs from Python callables.",
-    url="https://github.com/michaelosthege/aesara-federated",
-    download_url=f"https://github.com/michaelosthege/aesara-federated/tarball/{__version__}",
+    url="https://github.com/michaelosthege/pytensor-federated",
+    download_url=f"https://github.com/michaelosthege/pytensor-federated/tarball/{__version__}",
     author="Michael Osthege",
     author_email="[email protected]",
     license="GNU Affero General Public License v3",