pytensor 2.19.0__tar.gz → 2.21.0__tar.gz

{pytensor-2.19.0 → pytensor-2.21.0}/MANIFEST.in

@@ -8,6 +8,5 @@ global-include *.h
 global-include *.hh
 global-include *.sh
 recursive-include doc *
-include bin/pytensor-cache
 include versioneer.py
 include pytensor/_version.py

{pytensor-2.19.0/pytensor.egg-info → pytensor-2.21.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor
-Version: 2.19.0
+Version: 2.21.0
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License: .. _license:
@@ -66,16 +66,15 @@ Classifier: Operating System :: POSIX
 Classifier: Operating System :: Unix
 Classifier: Operating System :: MacOS
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Python: <3.13,>=3.9
+Requires-Python: <3.13,>=3.10
 Description-Content-Type: text/x-rst
 License-File: LICENSE.txt
 Requires-Dist: setuptools>=48.0.0
 Requires-Dist: scipy>=0.14
-Requires-Dist: numpy>=1.17.0
+Requires-Dist: numpy<2,>=1.17.0
 Requires-Dist: filelock
 Requires-Dist: etuples
 Requires-Dist: logical-unification
@@ -106,8 +105,7 @@ Provides-Extra: jax
 Requires-Dist: jax; extra == "jax"
 Requires-Dist: jaxlib; extra == "jax"
 Provides-Extra: numba
-Requires-Dist: numba>=0.55; extra == "numba"
-Requires-Dist: numba-scipy>=0.3.0; extra == "numba"
+Requires-Dist: numba>=0.57; extra == "numba"
 
 .. image:: https://cdn.rawgit.com/pymc-devs/pytensor/main/doc/images/PyTensor_RGB.svg
     :height: 100px

{pytensor-2.19.0 → pytensor-2.21.0}/doc/conf.py

@@ -121,6 +121,32 @@ html_logo = "images/PyTensor_RGB.svg"
 html_theme = "pymc_sphinx_theme"
 html_theme_options = {
     "use_search_override": False,
+    "icon_links": [
+        {
+            "url": "https://github.com/pymc-devs/pytensor/",
+            "icon": "fa-brands fa-github",
+            "name": "GitHub",
+            "type": "fontawesome",
+        },
+        {
+            "url": "https://twitter.com/pymc_devs/",
+            "icon": "fa-brands fa-twitter",
+            "name": "Twitter",
+            "type": "fontawesome",
+        },
+        {
+            "url": "https://www.youtube.com/c/PyMCDevelopers",
+            "icon": "fa-brands fa-youtube",
+            "name": "YouTube",
+            "type": "fontawesome",
+        },
+        {
+            "url": "https://discourse.pymc.io",
+            "icon": "fa-brands fa-discourse",
+            "name": "Discourse",
+            "type": "fontawesome",
+        },
+    ],
 }
 html_context = {
     "github_user": "pymc-devs",
@@ -190,6 +216,7 @@ html_use_smartypants = True
 # Output file base name for HTML help builder.
 htmlhelp_basename = "pytensor_doc"
 
+
 # Options for the linkcode extension
 # ----------------------------------
 # Resolve function
@@ -239,7 +266,13 @@ latex_elements = {
 # (source start file, target name, title, author, document class
 # [howto/manual]).
 latex_documents = [
-    ("index", "pytensor.tex", "PyTensor Documentation", "PyTensor Developers", "manual"),
+    (
+        "index",
+        "pytensor.tex",
+        "PyTensor Documentation",
+        "PyTensor Developers",
+        "manual",
+    ),
 ]
 
 # The name of an image file (relative to this directory) to place at the top of

{pytensor-2.19.0 → pytensor-2.21.0}/doc/dev_start_guide.rst

@@ -168,15 +168,9 @@ create a virtual environment in the project directory:
 
 .. code-block:: bash
 
-    conda env create -n pytensor-dev -f environment.yml
+    conda env create -f environment.yml
     conda activate pytensor-dev
 
-Afterward, you can install the development dependencies:
-
-.. code-block:: bash
-
-    pip install -r requirements.txt
-
 Next, ``pre-commit`` needs to be configured so that the linting and code quality
 checks are performed before each commit:
 
@@ -202,12 +196,12 @@ see the `NumPy development guide <https://numpy.org/doc/stable/dev/>`_.
 Contributing to the documentation
 ---------------------------------
 
-To contribute to the documentation, first follow the instructions in the previous section. Afterward, you can install the documentation dependencies in the virtual environment you created:
-
+The documentation build dependencies have also been included in the virtual environment you created. You can also create a separate virtual environment just for the documentation using the `environment.yml` file located inside the `doc` folder.
 
 .. code-block:: bash
 
-    pip install -r requirements-rtd.txt
+    conda env create -f doc/environment.yml
+    conda activate pytensor-docs
 
 
 You can now build the documentation from the root of the project with:
@@ -215,7 +209,7 @@ You can now build the documentation from the root of the project with:
 
 .. code-block:: bash
 
-    python doc/scripts/docgen.py
+    python -m sphinx -b html ./doc ./html
 
 
 Afterward, you can go to `html/index.html` and navigate the changes in a browser. One way to do this is to go to the `html` directory and run:
@@ -226,7 +220,7 @@ Afterward, you can go to `html/index.html` and navigate the changes in a browser
     python -m http.server
 
 **Do not commit the `html` directory. The documentation is built automatically.**
-
+For more documentation customizations such as different formats e.g., PDF, refer to the `Sphinx documentation <https://www.sphinx-doc.org/en/master/usage/builders/index.html>`_.
 
 Other tools that might help
 ===========================

{pytensor-2.19.0 → pytensor-2.21.0}/doc/environment.yml

@@ -3,7 +3,7 @@ channels:
   - conda-forge
   - nodefaults
 dependencies:
-  - python=3.9
+  - python=3.10
   - gcc_linux-64
   - gxx_linux-64
   - numpy
@@ -12,7 +12,7 @@ dependencies:
   - sphinx>=5.1.0,<6
   - mock
   - pillow
+  - pymc-sphinx-theme
   - pip
   - pip:
-    - git+https://github.com/pymc-devs/pymc-sphinx-theme
     - -e ..

pytensor-2.21.0/doc/extending/creating_a_numba_jax_op.rst

@@ -0,0 +1,285 @@
+Adding JAX and Numba support for `Op`\s
+=======================================
+
+PyTensor is able to convert its graphs into JAX and Numba compiled functions. In order to do
+this, each :class:`Op` in an PyTensor graph must have an equivalent JAX/Numba implementation function.
+
+This tutorial will explain how JAX and Numba implementations are created for an :class:`Op`.  It will
+focus specifically on the JAX case, but the same mechanisms are used for Numba as well.
+
+Step 1: Identify the PyTensor :class:`Op` you'd like to implement in JAX
+------------------------------------------------------------------------
+
+Find the source for the PyTensor :class:`Op` you'd like to be supported in JAX, and
+identify the function signature and return values. These can be determined by
+looking at the :meth:`Op.make_node` implementation. In general, one needs to be familiar
+with PyTensor :class:`Op`\s in order to provide a conversion implementation, so first read
+:ref:`creating_an_op` if you are not familiar.
+
+For example, you want to extend support for :class:`CumsumOp`\:
+
+.. code:: python
+
+    class CumsumOp(Op):
+        __props__ = ("axis",)
+
+        def __new__(typ, *args, **kwargs):
+            obj = object.__new__(CumOp, *args, **kwargs)
+            obj.mode = "add"
+            return obj
+
+
+:class:`CumsumOp` turns out to be a variant of :class:`CumOp`\ :class:`Op`
+which currently has an :meth:`Op.make_node` as follows:
+
+.. code:: python
+
+    def make_node(self, x):
+        x = ptb.as_tensor_variable(x)
+        out_type = x.type()
+
+        if self.axis is None:
+            out_type = vector(dtype=x.dtype)  # Flatten
+        elif self.axis >= x.ndim or self.axis < -x.ndim:
+            raise ValueError(f"axis(={self.axis}) out of bounds")
+
+        return Apply(self, [x], [out_type])
+
+The :class:`Apply` instance that's returned specifies the exact types of inputs that
+our JAX implementation will receive and the exact types of outputs it's expected to
+return--both in terms of their data types and number of dimensions/shapes.
+The actual inputs our implementation will receive are necessarily numeric values
+or NumPy :class:`ndarray`\s; all that :meth:`Op.make_node` tells us is the
+general signature of the underlying computation.
+
+More specifically, the :class:`Apply` implies that there is one input that is
+automatically converted to PyTensor variables via :func:`as_tensor_variable`.
+There is another parameter, `axis`, that is used to determine the direction
+of the operation, hence shape of the output. The check that follows imply that
+`axis` must refer to a dimension in the input tensor. The input's elements
+could also have any data type (e.g. floats, ints), so our JAX implementation
+must be able to handle all the possible data types.
+
+It also tells us that there's only one return value, that it has a data type
+determined by :meth:`x.type()` i.e., the data type of the original tensor.
+This implies that the result is necessarily a matrix.
+
+Some class may have a more complex behavior. For example, the :class:`CumOp`\ :class:`Op`
+also has another variant :class:`CumprodOp`\ :class:`Op` with the exact signature
+as :class:`CumsumOp`\ :class:`Op`. The difference lies in that the `mode` attribute in
+:class:`CumOp` definition:
+
+.. code:: python
+
+    class CumOp(COp):
+        # See function cumsum/cumprod for docstring
+
+        __props__ = ("axis", "mode")
+        check_input = False
+        params_type = ParamsType(
+            c_axis=int_t, mode=EnumList(("MODE_ADD", "add"), ("MODE_MUL", "mul"))
+        )
+
+        def __init__(self, axis: int | None = None, mode="add"):
+            if mode not in ("add", "mul"):
+                raise ValueError(f'{type(self).__name__}: Unknown mode "{mode}"')
+            self.axis = axis
+            self.mode = mode
+
+        c_axis = property(lambda self: np.MAXDIMS if self.axis is None else self.axis)
+
+`__props__` is used to parametrize the general behavior of the :class:`Op`. One need to
+pay attention to this to decide whether the JAX implementation should support all variants
+or raise an explicit NotImplementedError for cases that are not supported e.g., when
+:class:`CumsumOp` of :class:`CumOp("add")` is supported but not :class:`CumprodOp` of
+:class:`CumOp("mul")`.
+
+Next, we look at the :meth:`Op.perform` implementation to see exactly
+how the inputs and outputs are used to compute the outputs for an :class:`Op`
+in Python. This method is effectively what needs to be implemented in JAX.
+
+Step 2: Find the relevant JAX method (or something close)
+---------------------------------------------------------
+
+With a precise idea of what the PyTensor :class:`Op` does we need to figure out how
+to implement it in JAX. In the best case scenario, JAX has a similarly named
+function that performs exactly the same computations as the :class:`Op`. For
+example, the :class:`Eye` operator has a JAX equivalent: :func:`jax.numpy.eye`
+(see `the documentation <https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.eye.html?highlight=eye>`_).
+
+If we wanted to implement an :class:`Op` like :class:`IfElse`, we might need to
+recreate the functionality with some custom logic.  In many cases, at least some
+custom logic is needed to reformat the inputs and outputs so that they exactly
+match the `Op`'s.
+
+Here's an example for :class:`IfElse`:
+
+.. code:: python
+
+   def ifelse(cond, *args, n_outs=n_outs):
+       res = jax.lax.cond(
+           cond, lambda _: args[:n_outs], lambda _: args[n_outs:], operand=None
+       )
+       return res if n_outs > 1 else res[0]
+
+In this case, :class:`CumOp` is implemented with NumPy's :func:`numpy.cumsum`
+and :func:`numpy.cumprod`, which have JAX equivalents: :func:`jax.numpy.cumsum`
+and :func:`jax.numpy.cumprod`.
+
+.. code:: python
+
+    def perform(self, node, inputs, output_storage):
+        x = inputs[0]
+        z = output_storage[0]
+        if self.mode == "add":
+            z[0] = np.cumsum(x, axis=self.axis)
+        else:
+            z[0] = np.cumprod(x, axis=self.axis)
+
+Step 3: Register the function with the `jax_funcify` dispatcher
+---------------------------------------------------------------
+
+With the PyTensor `Op` replicated in JAX, we'll need to register the
+function with the PyTensor JAX `Linker`. This is done through the use of
+`singledispatch`. If you don't know how `singledispatch` works, see the
+`Python documentation <https://docs.python.org/3/library/functools.html#functools.singledispatch>`_.
+
+The relevant dispatch functions created by `singledispatch` are :func:`pytensor.link.numba.dispatch.numba_funcify` and
+:func:`pytensor.link.jax.dispatch.jax_funcify`.
+
+Here's an example for the `CumOp`\ `Op`:
+
+.. code:: python
+
+   import jax.numpy as jnp
+
+   from pytensor.tensor.extra_ops import CumOp
+   from pytensor.link.jax.dispatch import jax_funcify
+
+
+    @jax_funcify.register(CumOp)
+    def jax_funcify_CumOp(op, **kwargs):
+        axis = op.axis
+        mode = op.mode
+
+        def cumop(x, axis=axis, mode=mode):
+            if mode == "add":
+                return jnp.cumsum(x, axis=axis)
+            else:
+                return jnp.cumprod(x, axis=axis)
+
+        return cumop
+
+Suppose `jnp.cumprod` does not exist, we will need to register the function as follows:
+
+.. code:: python
+
+   import jax.numpy as jnp
+
+   from pytensor.tensor.extra_ops import CumOp
+   from pytensor.link.jax.dispatch import jax_funcify
+
+
+    @jax_funcify.register(CumOp)
+    def jax_funcify_CumOp(op, **kwargs):
+        axis = op.axis
+        mode = op.mode
+
+        def cumop(x, axis=axis, mode=mode):
+            if mode == "add":
+                return jnp.cumsum(x, axis=axis)
+            else:
+                raise NotImplementedError("JAX does not support cumprod function at the moment.")
+
+        return cumop
+
+Step 4: Write tests
+-------------------
+
+Test that your registered `Op` is working correctly by adding tests to the
+appropriate test suites in PyTensor (e.g. in ``tests.link.jax`` and one of
+the modules in ``tests.link.numba``). The tests should ensure that your implementation can
+handle the appropriate types of inputs and produce outputs equivalent to `Op.perform`.
+Check the existing tests for the general outline of these kinds of tests. In
+most cases, a helper function can be used to easily verify the correspondence
+between a JAX/Numba implementation and its `Op`.
+
+For example, the :func:`compare_jax_and_py` function streamlines the steps
+involved in making comparisons with `Op.perform`.
+
+Here's a small example of a test for :class:`CumOp` above:
+
+.. code:: python
+    
+    import numpy as np
+    import pytensor.tensor as pt
+    from pytensor.configdefaults import config
+    from tests.link.jax.test_basic import compare_jax_and_py
+    from pytensor.graph import FunctionGraph
+    from pytensor.graph.op import get_test_value
+
+    def test_jax_CumOp():
+        """Test JAX conversion of the `CumOp` `Op`."""
+
+        # Create a symbolic input for the first input of `CumOp`
+        a = pt.matrix("a")
+
+        # Create test value tag for a
+        a.tag.test_value = np.arange(9, dtype=config.floatX).reshape((3, 3))
+
+        # Create the output variable
+        out = pt.cumsum(a, axis=0)
+
+        # Create a PyTensor `FunctionGraph`
+        fgraph = FunctionGraph([a], [out])
+
+        # Pass the graph and inputs to the testing function
+        compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
+
+        # For the second mode of CumOp
+        out = pt.cumprod(a, axis=1)
+        fgraph = FunctionGraph([a], [out])
+        compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
+
+If the variant :class:`CumprodOp` is not implemented, we can add a test for it as follows:
+
+.. code:: python
+
+    import pytest
+    
+    def test_jax_CumOp():
+        """Test JAX conversion of the `CumOp` `Op`."""
+        a = pt.matrix("a")
+        a.tag.test_value = np.arange(9, dtype=config.floatX).reshape((3, 3))
+        
+        with pytest.raises(NotImplementedError):
+            out = pt.cumprod(a, axis=1)
+            fgraph = FunctionGraph([a], [out])
+            compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
+
+Note
+----
+In out previous example of extending JAX, :class:`Eye`\ :class:`Op` was used with the test function as follows:
+
+.. code:: python
+    def test_jax_Eye():
+        """Test JAX conversion of the `Eye` `Op`."""
+
+        # Create a symbolic input for `Eye`
+        x_at = pt.scalar()
+
+        # Create a variable that is the output of an `Eye` `Op`
+        eye_var = pt.eye(x_at)
+
+        # Create an PyTensor `FunctionGraph`
+        out_fg = FunctionGraph(outputs=[eye_var])
+
+        # Pass the graph and any inputs to the testing function
+        compare_jax_and_py(out_fg, [3])
+
+This one nowadays leads to a test failure due to new restrictions in JAX + JIT,
+as reported in issue `#654 <https://github.com/pymc-devs/pytensor/issues/654>`_.
+All jitted functions now must have constant shape, which means a graph like the
+one of :class:`Eye` can never be translated to JAX, since it's fundamentally a
+function with dynamic shapes. In other words, only PyTensor graphs with static shapes
+can be translated to JAX at the moment.

{pytensor-2.19.0 → pytensor-2.21.0}/doc/extending/pipeline.rst

@@ -88,7 +88,7 @@ case if ``borrow`` was True, the thunk would be allowed to reuse--or
     The compile cache is based upon the C++ code of the graph to be compiled.
     So, if you change compilation configuration variables, such as
     :attr:`config.blas__ldflags`, you will need to manually remove your compile cache,
-    using ``PyTensor/bin/pytensor-cache clear``
+    using ``pytensor-cache clear``
 
     PyTensor also implements a lock mechanism that prevents multiple compilations
     within the same compilation directory (to avoid crashes with parallel

{pytensor-2.19.0 → pytensor-2.21.0}/pyproject.toml

@@ -10,7 +10,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "pytensor"
 dynamic = ['version']
-requires-python = ">=3.9,<3.13"
+requires-python = ">=3.10,<3.13"
 authors = [{ name = "pymc-devs", email = "pymc.devs@gmail.com" }]
 description = "Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs."
 readme = "README.rst"
@@ -30,7 +30,6 @@ classifiers = [
     "Operating System :: Unix",
     "Operating System :: MacOS",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
@@ -49,7 +48,7 @@ keywords = [
 dependencies = [
     "setuptools>=48.0.0",
     "scipy>=0.14",
-    "numpy>=1.17.0",
+    "numpy>=1.17.0,<2",
     "filelock",
     "etuples",
     "logical-unification",
@@ -80,10 +79,10 @@ tests = [
 ]
 rtd = ["sphinx>=5.1.0,<6", "pygments", "pydot", "pydot2", "pydot-ng"]
 jax = ["jax", "jaxlib"]
-numba = ["numba>=0.55", "numba-scipy>=0.3.0"]
+numba = ["numba>=0.57"]
 
 [tool.setuptools.packages.find]
-include = ["pytensor*", "bin"]
+include = ["pytensor*"]
 
 [tool.setuptools.package-data]
 pytensor = ["py.typed"]
@@ -123,7 +122,7 @@ testpaths = "tests/"
 
 [tool.ruff]
 line-length = 88
-exclude = ["doc/", "pytensor/_version.py", "bin/pytensor_cache.py"]
+exclude = ["doc/", "pytensor/_version.py"]
 
 [tool.ruff.lint]
 select = ["C", "E", "F", "I", "UP", "W", "RUF"]
@@ -154,6 +153,7 @@ lines-after-imports = 2
 
 
 [tool.mypy]
+python_version = "3.10"
 ignore_missing_imports = true
 no_implicit_optional = true
 check_untyped_defs = false

{pytensor-2.19.0 → pytensor-2.21.0}/pytensor/__init__.py

@@ -18,7 +18,6 @@ To learn more, check out:
 
 """
 
-
 __docformat__ = "restructuredtext en"
 
 # Set a default logger. It is important to do this before importing some other
@@ -77,7 +76,7 @@ from pytensor.graph.replace import clone_replace, graph_replace
 # isort: on
 
 
-def as_symbolic(x: Any, name: Optional[str] = None, **kwargs) -> Variable:
+def as_symbolic(x: Any, name: str | None = None, **kwargs) -> Variable:
     """Convert `x` into an equivalent PyTensor `Variable`.
 
     Parameters

{pytensor-2.19.0 → pytensor-2.21.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-03-08T17:58:24+0000",
+ "date": "2024-05-13T21:55:31+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "4ee358818543ee271b7922f5ccf988ef9338591f",
- "version": "2.19.0"
+ "full-revisionid": "90da9e61688b0fef8b91106f0e2ae08ee67be1be",
+ "version": "2.21.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.19.0 → pytensor-2.21.0}/pytensor/bin/pytensor_cache.py

@@ -75,9 +75,9 @@ def main():
             if items:
                 _logger.warning(
                     "There remain elements in the cache dir that you may "
-                    "need to erase manually. The cache dir is:\n  %s\n"
+                    f"need to erase manually. The cache dir is:\n  {config.compiledir}\n"
                     'You can also call "pytensor-cache purge" to '
-                    "remove everything from that directory." % config.compiledir
+                    "remove everything from that directory."
                 )
                 _logger.debug(f"Remaining elements ({len(items)}): {', '.join(items)}")
         elif sys.argv[1] == "list":

{pytensor-2.19.0 → pytensor-2.21.0}/pytensor/breakpoint.py

@@ -105,7 +105,7 @@ class PdbBreakpoint(Op):
             except Exception:
                 raise ValueError(
                     "Some of the inputs to the PdbBreakpoint op "
-                    "'%s' could not be casted to NumPy arrays" % self.name
+                    f"'{self.name}' could not be casted to NumPy arrays"
                 )
 
             print("\n")