pytensor 2.20.0__tar.gz → 2.22.0__tar.gz

{pytensor-2.20.0 → pytensor-2.22.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.20.0/pytensor.egg-info → pytensor-2.22.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor
-Version: 2.20.0
+Version: 2.22.0
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License: .. _license:
@@ -74,7 +74,7 @@ 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

{pytensor-2.20.0 → pytensor-2.22.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.20.0 → pytensor-2.22.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.20.0 → pytensor-2.22.0}/doc/environment.yml

@@ -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.22.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.20.0 → pytensor-2.22.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.20.0 → pytensor-2.22.0}/pyproject.toml

@@ -48,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",
@@ -82,7 +82,7 @@ jax = ["jax", "jaxlib"]
 numba = ["numba>=0.57"]
 
 [tool.setuptools.packages.find]
-include = ["pytensor*", "bin"]
+include = ["pytensor*"]
 
 [tool.setuptools.package-data]
 pytensor = ["py.typed"]
@@ -122,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"]

{pytensor-2.20.0 → pytensor-2.22.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-03-28T18:19:27+0100",
+ "date": "2024-05-22T14:12:17+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "ef22377d7e6ea6c74c36a8dc634ef7f1f1f8f9c6",
- "version": "2.20.0"
+ "full-revisionid": "bb028ae2330433755b9d4aa32ab6e8d0c9f662fc",
+ "version": "2.22.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.20.0 → pytensor-2.22.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.20.0 → pytensor-2.22.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")

{pytensor-2.20.0 → pytensor-2.22.0}/pytensor/compile/builders.py

@@ -7,7 +7,7 @@ from functools import partial
 from typing import cast
 
 import pytensor.tensor as pt
-from pytensor import function
+from pytensor.compile.function import function
 from pytensor.compile.function.pfunc import rebuild_collect_shared
 from pytensor.compile.mode import optdb
 from pytensor.compile.sharedvalue import SharedVariable
@@ -400,6 +400,15 @@ class OpFromGraph(Op, HasInnerGraph):
             Check :func:`pytensor.function` for more arguments, only works when not
             inline.
         """
+        ignore_unused_inputs = kwargs.get("on_unused_input", False) == "ignore"
+        if not ignore_unused_inputs and len(inputs) != len(set(inputs)):
+            var_counts = {var: inputs.count(var) for var in inputs}
+            duplicated_inputs = [var for var, count in var_counts.items() if count > 1]
+            raise ValueError(
+                f"There following variables were provided more than once as inputs to the OpFromGraph, resulting in an "
+                f"invalid graph: {duplicated_inputs}. Use dummy variables or var.copy() to distinguish "
+                f"variables when creating the OpFromGraph graph."
+            )
 
         if not (isinstance(inputs, list) and isinstance(outputs, list)):
             raise TypeError("Inputs and outputs must be lists")
@@ -700,7 +709,7 @@ class OpFromGraph(Op, HasInnerGraph):
             if not isinstance(roverrides_l, list):
                 raise TypeError(
                     "Rop overriding function should return a list, "
-                    'got "%s"' % type(roverrides_l)
+                    f'got "{type(roverrides_l)}"'
                 )
             all_rops_l, all_rops_ov_l = zip(
                 *[

{pytensor-2.20.0 → pytensor-2.22.0}/pytensor/compile/mode.py

@@ -452,13 +452,11 @@ JAX = Mode(
     JAXLinker(),
     RewriteDatabaseQuery(
         include=["fast_run", "jax"],
-        # TODO: "local_uint_constant_indices" can be reintroduced once https://github.com/google/jax/issues/16836 is fixed.
         exclude=[
             "cxx_only",
             "BlasOpt",
             "fusion",
             "inplace",
-            "local_uint_constant_indices",
         ],
     ),
 )

{pytensor-2.20.0 → pytensor-2.22.0}/pytensor/compile/ops.py

@@ -252,7 +252,7 @@ class FromFunctionOp(Op):
         return hash(type(self)) ^ hash(self.__fn)
 
     def __str__(self):
-        return "FromFunctionOp{%s}" % self.__fn.__name__
+        return f"FromFunctionOp{{{self.__fn.__name__}}}"
 
     def perform(self, node, inputs, outputs):
         outs = self.__fn(*inputs)

{pytensor-2.20.0 → pytensor-2.22.0}/pytensor/configdefaults.py

@@ -1437,12 +1437,16 @@ add_vm_configvars()
 add_numba_configvars()
 
 # TODO: `gcc_version_str` is used by other modules.. Should it become an immutable config var?
-try:
-    p_out = output_subprocess_Popen([config.cxx, "-dumpversion"])
-    gcc_version_str = p_out[0].strip().decode()
-except OSError:
-    # Typically means gcc cannot be found.
+if config.cxx != "":
+    try:
+        p_out = output_subprocess_Popen([config.cxx, "-dumpversion"])
+        gcc_version_str = p_out[0].strip().decode()
+    except OSError:
+        # Typically means gcc cannot be found.
+        gcc_version_str = "GCC_NOT_FOUND"
+else:
     gcc_version_str = "GCC_NOT_FOUND"
+
 # TODO: The caching dir resolution is a procedural mess of helper functions, local variables
 # and config definitions. And the result is also not particularly pretty..
 add_caching_dir_configvars()