Add support for PyMC >= 5 with PyTensor backend

.github/workflows/tests.yml

@@ -17,8 +17,8 @@ jobs:
           - "test"
           - "test_comparison"
           - "test_pymc3"
-          - "test_pymc4"
-          - "test_pymc4_jax"
+          - "test_pymc"
+          - "test_pymc_jax"
           - "test_jax"
 
     steps:

.github/workflows/wheels.yml

@@ -14,7 +14,7 @@ jobs:
         with:
           submodules: true
           fetch-depth: 0
-      - uses: pypa/cibuildwheel@v2.15.0
+      - uses: pypa/cibuildwheel@v2.16.2
         env:
           CIBW_SKIP: "pp* *-win32 *-manylinux_i686 *-musllinux*"
           CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.4.0
+    rev: v4.5.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer

noxfile.py

@@ -1,6 +1,6 @@
 import nox
 
-ALL_PYTHON_VS = ["3.8", "3.9", "3.10"]
+ALL_PYTHON_VS = ["3.9", "3.10", "3.11"]
 
 
 @nox.session(python=ALL_PYTHON_VS)
@@ -24,20 +24,20 @@ def test_pymc3(session):
 
 
 @nox.session(python=ALL_PYTHON_VS)
-def test_pymc4(session):
-    session.install(".[test,pymc4]")
-    session.run("python", "-c", "import aesara")
-    session.run("python", "-c", "import exoplanet_core.pymc4.ops")
-    session.run("pytest", "-v", "tests/pymc4_test.py", *session.posargs)
+def test_pymc(session):
+    session.install(".[test,pymc]")
+    session.run("python", "-c", "import pytensor")
+    session.run("python", "-c", "import exoplanet_core.pymc.ops")
+    session.run("pytest", "-v", "tests/pymc_test.py", *session.posargs)
 
 
 @nox.session(python=ALL_PYTHON_VS)
-def test_pymc4_jax(session):
-    session.install(".[test,pymc4,jax]")
+def test_pymc_jax(session):
+    session.install(".[test,pymc,jax]")
     session.run("python", "-c", "import jax")
-    session.run("python", "-c", "import aesara")
-    session.run("python", "-c", "import exoplanet_core.pymc4.ops")
-    session.run("pytest", "-v", "tests/pymc4_jax_test.py", *session.posargs)
+    session.run("python", "-c", "import pytensor")
+    session.run("python", "-c", "import exoplanet_core.pymc.ops")
+    session.run("pytest", "-v", "tests/pymc_jax_test.py", *session.posargs)
 
 
 @nox.session(python=ALL_PYTHON_VS)
@@ -50,7 +50,7 @@ def test_jax(session):
 
 @nox.session(python=ALL_PYTHON_VS)
 def test_all(session):
-    session.install(".[test,pymc3,pymc4,jax,comparison]")
+    session.install(".[test,pymc3,pymc,jax,comparison]")
     session.run("python", "-c", "import jax")
     session.run("python", "-c", "import aesara")
     session.run("python", "-c", "import theano")

setup.py

@@ -31,10 +31,15 @@
 INSTALL_REQUIRES = ["numpy>=1.13.0"]
 EXTRA_REQUIRE = {
     "pymc3": ["pymc3>=3.9", "numpy<1.22"],
-    "pymc4": ["pymc>=4.0.0,<5"],
+    "pymc": ["pymc>=5.0.0"],
     "jax": ["jax", "jaxlib"],
     "test": ["pytest"],
-    "comparison": ["batman-package", "starry", "numpy<1.22"],
+    "comparison": [
+        "batman-package",
+        "starry",
+        "numpy<1.22",
+        "xarray<2023.10.0",
+    ],
     "benchmark": [
         "pytest",
         "pytest-benchmark",

src/exoplanet_core/jax/ops.py

@@ -7,7 +7,7 @@
 import numpy as np
 from jax import core
 from jax import numpy as jnp
-from jax.abstract_arrays import ShapedArray
+from jax.core import ShapedArray
 from jax.interpreters import ad, batching, xla
 from jax.lib import xla_client
 
@@ -16,7 +16,7 @@
 xops = xla_client.ops
 
 for _name, _value in cpu_driver.registrations().items():
-    xla_client.register_cpu_custom_call_target(_name, _value)
+    xla_client.register_custom_call_target(_name, _value, platform="cpu")
 
 
 try:

src/exoplanet_core/pymc4/__init__.py → src/exoplanet_core/pymc/__init__.py

@@ -4,27 +4,27 @@
 
 
 def __set_compiler_flags():
-    import aesara
+    import pytensor
 
     def add_flag(current, new):
         if new in current:
             return current
         return f"{current} {new}"
 
-    current = aesara.config.gcc__cxxflags
+    current = pytensor.config.gcc__cxxflags
     current = add_flag(current, "-Wno-c++11-narrowing")
     current = add_flag(current, "-fno-exceptions")
     current = add_flag(current, "-fno-unwind-tables")
     current = add_flag(current, "-fno-asynchronous-unwind-tables")
-    aesara.config.gcc__cxxflags = current
+    pytensor.config.gcc__cxxflags = current
 
 
 __set_compiler_flags()
 
 
-from exoplanet_core.pymc4 import ops
+from exoplanet_core.pymc import ops
 
 try:
-    from exoplanet_core.pymc4 import jax_support  # noqa
+    from exoplanet_core.pymc import jax_support  # noqa
 except ImportError:
     pass

src/exoplanet_core/pymc4/jax_support.py → src/exoplanet_core/pymc/jax_support.py

@@ -1,26 +1,26 @@
-from aesara.link.jax.dispatch import jax_funcify
+from pytensor.link.jax.dispatch import jax_funcify
 
 from exoplanet_core.jax import ops as jax_ops
-from exoplanet_core.pymc4 import ops as pymc4_ops
+from exoplanet_core.pymc import ops as pymc_ops
 
 
-@jax_funcify.register(pymc4_ops.Kepler)
+@jax_funcify.register(pymc_ops.Kepler)
 def jax_funcify_Kepler(op, **kwargs):
     def kepler(M, ecc):
         return jax_ops.kepler(M, ecc)
 
     return kepler
 
 
-@jax_funcify.register(pymc4_ops.QuadSolutionVector)
+@jax_funcify.register(pymc_ops.QuadSolutionVector)
 def jax_funcify_QuadSolutionVector(op, **kwargs):
     def quad_solution_vector(b, r):
         return jax_ops._base_quad_solution_vector(b, r)
 
     return quad_solution_vector
 
 
-@jax_funcify.register(pymc4_ops.ContactPoints)
+@jax_funcify.register(pymc_ops.ContactPoints)
 def jax_funcify_ContactPoints(op, **kwargs):
     def contact_points(a, e, cosw, sinw, cosi, sini, L):
         return jax_ops.contact_points(a, e, cosw, sinw, cosi, sini, L)

src/exoplanet_core/pymc4/ops.py → src/exoplanet_core/pymc/ops.py

@@ -4,16 +4,16 @@
 
 from itertools import chain
 
-import aesara
-import aesara.tensor as at
 import numpy as np
-from aesara.graph import basic, op
+import pytensor
+import pytensor.tensor as pt
+from pytensor.graph import basic, op
 
 from exoplanet_core import driver
 
 
 def as_tensor_variable(x, dtype="float64", **kwargs):
-    t = at.as_tensor_variable(x, **kwargs)
+    t = pt.as_tensor_variable(x, **kwargs)
     if dtype is None:
         return t
     return t.astype(dtype)
@@ -84,13 +84,17 @@ def grad(self, inputs, gradients):
         dfdM = (1 + ecosf) ** 2 / ome2**1.5
         dfde = (2 + ecosf) * sinf / ome2
 
-        bM = at.zeros_like(M)
-        be = at.zeros_like(M)
-        if not isinstance(gradients[0].type, aesara.gradient.DisconnectedType):
+        bM = pt.zeros_like(M)
+        be = pt.zeros_like(M)
+        if not isinstance(
+            gradients[0].type, pytensor.gradient.DisconnectedType
+        ):
             bM += gradients[0] * cosf * dfdM
             be += gradients[0] * cosf * dfde
 
-        if not isinstance(gradients[1].type, aesara.gradient.DisconnectedType):
+        if not isinstance(
+            gradients[1].type, pytensor.gradient.DisconnectedType
+        ):
             bM -= gradients[1] * sinf * dfdM
             be -= gradients[1] * sinf * dfde
 
@@ -137,8 +141,13 @@ def make_node(self, b, r):
             raise ValueError("float64 precision is required")
         x = in_args[0]
         o = [
-            at.tensor(
-                broadcastable=tuple(x.broadcastable) + (False,),
+            # NOTE: Changed arg broadcastable to shape because of deprecation warning,
+            # BUT this might change again in the future: https://github.com/pymc-devs/pytensor/issues/408
+            # The other option is to use `x.type.shape + (3,)`
+            # Ref: https://pytensor.readthedocs.io/en/latest/library/tensor/basic.html#pytensor.tensor.TensorType
+            pt.tensor(
+                # PyTensor internally converts True to 1 and False to None
+                shape=tuple(x.broadcastable) + (False,),
                 dtype=x.dtype,
             )
             for _ in range(3)
@@ -164,18 +173,18 @@ def grad(self, inputs, gradients):
         bs = gradients[0]
 
         for g in gradients[1:]:
-            if not isinstance(g.type, aesara.gradient.DisconnectedType):
+            if not isinstance(g.type, pytensor.gradient.DisconnectedType):
                 raise ValueError(
                     "Backpropagation is only supported for the solution vector"
                 )
 
-        if isinstance(bs.type, aesara.gradient.DisconnectedType):
+        if isinstance(bs.type, pytensor.gradient.DisconnectedType):
             return [
-                aesara.gradient.DisconnectedType()(),
-                aesara.gradient.DisconnectedType()(),
+                pytensor.gradient.DisconnectedType()(),
+                pytensor.gradient.DisconnectedType()(),
             ]
 
-        return [at.sum(bs * dsdb, axis=-1), at.sum(bs * dsdr, axis=-1)]
+        return [pt.sum(bs * dsdb, axis=-1), pt.sum(bs * dsdr, axis=-1)]
 
     def R_op(self, inputs, eval_points):
         if eval_points[0] is None:
@@ -205,8 +214,13 @@ def make_node(self, a, e, cosw, sinw, cosi, sini, L):
         out_args = [
             in_args[0].type(),
             in_args[0].type(),
-            at.tensor(
-                broadcastable=tuple(in_args[0].broadcastable),
+            # NOTE: Changed arg broadcastable to shape because of deprecation warning,
+            # BUT this might change again in the future: https://github.com/pymc-devs/pytensor/issues/408
+            # The other option is to use `in_args[0].type.shape`
+            # Ref: https://pytensor.readthedocs.io/en/latest/library/tensor/basic.html#pytensor.tensor.TensorType
+            pt.tensor(
+                # PyTensor internally converts True to 1 and False to None
+                shape=tuple(in_args[0].broadcastable),
                 dtype="int32",
             ),
         ]