Merge branch 'main' into inplace-linalg

Author: jessegrabowski

.github/workflows/test.yml

@@ -44,6 +44,7 @@ jobs:
             - '.github/workflows/*.yml'
             - 'setup.cfg'
             - 'requirements.txt'
+            - '.pre-commit-config.yaml'
 
   style:
     name: Check code style
@@ -52,7 +53,7 @@ jobs:
     if: ${{ needs.changes.outputs.changes == 'true' }}
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11"]
+        python-version: ["3.9", "3.12"]
     steps:
     - uses: actions/checkout@v4
     - uses: actions/setup-python@v5
@@ -70,7 +71,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.9", "3.11"]
+        python-version: ["3.9", "3.12"]
         fast-compile: [0,1]
         float32: [0,1]
         install-numba: [0]
@@ -101,7 +102,7 @@ jobs:
             float32: 0
             part: "tests/link/numba"
           - install-numba: 1
-            python-version: "3.11"
+            python-version: "3.11"  # TODO: bump to 3.12 when numba 0.59 is released
             fast-compile: 0
             float32: 0
             part: "tests/link/numba"
@@ -111,7 +112,7 @@ jobs:
             float32: 0
             part: "tests/link/jax"
           - install-jax: 1
-            python-version: "3.11"
+            python-version: "3.12"
             fast-compile: 0
             float32: 0
             part: "tests/link/jax"
@@ -139,7 +140,7 @@ jobs:
       - name: Install dependencies
         shell: bash -l {0}
         run: |
-          mamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl "numpy<1.26" scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock sympy
+          mamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock sympy
           # numba-scipy downgrades the installed scipy to 1.7.3 in Python 3.9, but
           # not numpy, even though scipy 1.7 requires numpy<1.23. When installing
           # PyTensor next, pip installs a lower version of numpy via the PyPI.
@@ -249,7 +250,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: "3.11"
+          python-version: "3.12"
 
       - name: Install dependencies
         run: |

.pre-commit-config.yaml

@@ -7,7 +7,7 @@ exclude: |
     )$
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.4.0
+    rev: v4.5.0
     hooks:
       - id: debug-statements
         exclude: |
@@ -19,42 +19,15 @@ repos:
               pytensor/tensor/variable\.py|
           )$
       - id: check-merge-conflict
-  - repo: https://github.com/asottile/pyupgrade
-    rev: v3.3.1
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.1.13
     hooks:
-      - id: pyupgrade
-        args: [--py39-plus]
-  - repo: https://github.com/psf/black
-    rev: 23.1.0
-    hooks:
-      - id: black
-        language_version: python3
-  - repo: https://github.com/pycqa/flake8
-    rev: 6.0.0
-    hooks:
-      - id: flake8
-        additional_dependencies:
-          - flake8-comprehensions
-  - repo: https://github.com/pycqa/isort
-    rev: 5.12.0
-    hooks:
-      - id: isort
-  - repo: https://github.com/humitos/mirrors-autoflake.git
-    rev: v1.1
-    hooks:
-      - id: autoflake
-        exclude: |
-          (?x)^(
-              .*/?__init__\.py|
-              pytensor/graph/toolbox\.py|
-              pytensor/link/jax/jax_dispatch\.py|
-              pytensor/link/jax/jax_linker\.py|
-              pytensor/scalar/basic_scipy\.py|
-              pytensor/tensor/linalg\.py
-          )$
-        args: ['--in-place', '--remove-all-unused-imports', '--remove-unused-variable']
+      - id: ruff
+        args: ["--fix", "--show-source"]
+      - id: ruff-format
+        args: ["--line-length=88"]
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.0.0
+    rev: v1.8.0
     hooks:
       - id: mypy
         language: python

README.rst

@@ -5,8 +5,9 @@
 
 |Tests Status| |Coverage|
 
-|Project Name| is a fork of `Aesara <https://github.com/aesara-devs/aesara>`__ -- a Python library that allows one to define, optimize, and
+|Project Name| is a Python library that allows one to define, optimize, and
 efficiently evaluate mathematical expressions involving multi-dimensional arrays.
+It provides the computational backend for `PyMC <https://github.com/pymc-devs/pymc>`__.
 
 Features
 ========
@@ -15,7 +16,8 @@ Features
 - Extensible graph framework suitable for rapid development of custom operators and symbolic optimizations
 - Implements an extensible graph transpilation framework that currently provides
   compilation via C, `JAX <https://github.com/google/jax>`__, and `Numba <https://github.com/numba/numba>`__
-- Based on one of the most widely-used Python tensor libraries: `Theano <https://github.com/Theano/Theano>`__
+- Contrary to PyTorch and TensorFlow, PyTensor maintains a static graph which can be modified in-place to
+  allow for advanced optimizations
 
 Getting started
 ===============
@@ -113,6 +115,11 @@ The current development branch of |Project Name| can be installed from GitHub, a
     pip install git+https://github.com/pymc-devs/pytensor
 
 
+Background
+==========
+
+PyTensor is a fork of `Aesara <https://github.com/aesara-devs/aesara>`__, which is a fork of `Theano <https://github.com/Theano/Theano>`__.
+
 Contributing
 ============
 

bin/__init__.py

@@ -1,5 +1,6 @@
 import warnings
 
+
 warnings.warn(
     message= "Importing 'bin.pytensor_cache' is deprecated. Import from "
     "'pytensor.bin.pytensor_cache' instead.",

environment.yml

@@ -39,12 +39,7 @@ dependencies:
   - pydot
   - ipython
   # code style
-  - black
-  - isort
-  # For linting
-  - flake8
-  - pep8
-  - pyflakes
+  - ruff
   # developer tools
   - pre-commit
   - packaging

pyproject.toml

@@ -1,24 +1,20 @@
 [build-system]
 requires = [
-  "setuptools>=48.0.0",
-  "cython",
-  "numpy>=1.17.0",
-  "versioneer[toml]==0.28",
+    "setuptools>=48.0.0",
+    "cython",
+    "numpy>=1.17.0",
+    "versioneer[toml]>=0.28",
 ]
 build-backend = "setuptools.build_meta"
 
 [project]
 name = "pytensor"
-dynamic = [
-  'version'
-]
-requires-python = ">=3.9,<3.12"
-authors = [
-    {name = "pymc-devs", email = "[email protected]"}
-]
+dynamic = ['version']
+requires-python = ">=3.9,<3.13"
+authors = [{ name = "pymc-devs", email = "[email protected]" }]
 description = "Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs."
 readme = "README.rst"
-license = {file = "LICENSE.txt"}
+license = { file = "LICENSE.txt" }
 classifiers = [
     "Development Status :: 6 - Mature",
     "Intended Audience :: Education",
@@ -37,6 +33,7 @@ classifiers = [
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
 ]
 
 keywords = [
@@ -71,15 +68,8 @@ documentation = "https://pytensor.readthedocs.io/en/latest/"
 pytensor-cache = "pytensor.bin.pytensor_cache:main"
 
 [project.optional-dependencies]
-complete = [
-    "pytensor[jax]",
-    "pytensor[numba]",
-]
-development = [
-    "pytensor[complete]",
-    "pytensor[tests]",
-    "pytensor[rtd]"
-]
+complete = ["pytensor[jax]", "pytensor[numba]"]
+development = ["pytensor[complete]", "pytensor[tests]", "pytensor[rtd]"]
 tests = [
     "pytest",
     "pre-commit",
@@ -88,34 +78,16 @@ tests = [
     "pytest-benchmark",
     "pytest-mock",
 ]
-rtd = [
-    "sphinx>=5.1.0,<6",
-    "pygments",
-    "pydot",
-    "pydot2",
-    "pydot-ng",
-]
-jax = [
-    "jax",
-    "jaxlib",
-]
-numba = [
-    "numba>=0.55",
-    "numba-scipy>=0.3.0"
-]
+rtd = ["sphinx>=5.1.0,<6", "pygments", "pydot", "pydot2", "pydot-ng"]
+jax = ["jax", "jaxlib"]
+numba = ["numba>=0.55", "numba-scipy>=0.3.0"]
 
 [tool.setuptools.packages.find]
 include = ["pytensor*", "bin"]
 
 [tool.setuptools.package-data]
-pytensor = [
-    "py.typed"
-]
-"pytensor.d3viz" = [
-    "html/*",
-    "css/*",
-    "js/*",
-]
+pytensor = ["py.typed"]
+"pytensor.d3viz" = ["html/*", "css/*", "js/*"]
 
 [tool.coverage.run]
 omit = [
@@ -136,14 +108,8 @@ branch = true
 relative_files = true
 
 [tool.coverage.report]
-omit = [
-    "pytensor/_version.py",
-    "tests/",
-]
-exclude_lines = [
-    "pragma: no cover",
-    "if TYPE_CHECKING:",
-]
+omit = ["pytensor/_version.py", "tests/"]
+exclude_lines = ["pragma: no cover", "if TYPE_CHECKING:"]
 show_missing = true
 
 [tool.versioneer]
@@ -163,41 +129,36 @@ max-line-length = 88
 [tool.pylint.messages_control]
 disable = ["C0330", "C0326"]
 
-[tool.isort]
-profile = "black"
-lines_after_imports = 2
-lines_between_sections = 1
-honor_noqa = true
-skip_gitignore = true
-skip = "pytensor/version.py"
-skip_glob = "**/*.pyx"
 
 [tool.ruff]
-select=["C","E","F","W"]
-ignore=["E501","E741","C408","C901"]
-exclude = [
-    "doc/",
-    "pytensor/_version.py",
-    "bin/pytensor_cache.py",
-]
+select = ["C", "E", "F", "I", "UP", "W"]
+ignore = ["C408", "C901", "E501", "E741", "UP031"]
+exclude = ["doc/", "pytensor/_version.py", "bin/pytensor_cache.py"]
+
+
+[tool.ruff.isort]
+lines-after-imports = 2
 
 [tool.ruff.per-file-ignores]
 # TODO: Get rid of these:
-    "**/__init__.py"=["F401","E402","F403"]
-    "pytensor/tensor/linalg.py"=["F401","F403"]
-    "pytensor/scalar/basic_scipy.py"=["E402","F403","F401"]
-    "pytensor/graph/toolbox.py"=["E402","F403","F401"]
-    "pytensor/link/jax/jax_dispatch.py"=["E402","F403","F401"]
-    "pytensor/link/jax/jax_linker.py"=["E402","F403","F401"]
-    "pytensor/sparse/sandbox/sp2.py"=["F401"]
-    "tests/tensor/test_math_scipy.py"=["E402"]
-    "tests/sparse/test_basic.py"=["E402"]
-    "tests/sparse/test_opt.py"=["E402"]
-    "tests/sparse/test_sp2.py"=["E402"]
-    "tests/sparse/test_utils.py"=["E402","F401"]
-    "tests/sparse/sandbox/test_sp.py"=["E402","F401"]
-    "tests/scalar/test_basic_sympy.py"=["E402"]
-    "pytensor/graph/rewriting/unify.py"=["F811"]
+"**/__init__.py" = ["F401", "E402", "F403"]
+"pytensor/tensor/linalg.py" = ["F401", "F403"]
+"pytensor/scalar/basic_scipy.py" = ["E402"]
+"pytensor/graph/toolbox.py" = ["E402"]
+# For the tests we skip because `pytest.importorskip` is used:
+"tests/link/jax/test_scalar.py" = ["E402"]
+"tests/link/jax/test_tensor_basic.py" = ["E402"]
+"tests/link/numba/test_basic.py" = ["E402"]
+"tests/link/numba/test_cython_support.py" = ["E402"]
+"tests/link/numba/test_performance.py" = ["E402"]
+"tests/link/numba/test_sparse.py" = ["E402"]
+"tests/link/numba/test_tensor_basic.py" = ["E402"]
+"tests/tensor/test_math_scipy.py" = ["E402"]
+"tests/sparse/test_basic.py" = ["E402"]
+"tests/sparse/test_sp2.py" = ["E402"]
+"tests/sparse/test_utils.py" = ["E402"]
+"tests/sparse/sandbox/test_sp.py" = ["E402", "F401"]
+"versioneer.py" = ["I"]
 
 
 [tool.mypy]

pytensor/compile/debugmode.py

@@ -1866,7 +1866,7 @@ def thunk():
                 # Nothing should be in storage map after evaluating
                 # each the thunk (specifically the last one)
                 for r, s in storage_map.items():
-                    assert type(s) is list
+                    assert isinstance(s, list)
                     assert s[0] is None
 
                 # store our output variables to their respective storage lists

pytensor/compile/profiling.py

@@ -1084,8 +1084,8 @@ def min_memory_generator(executable_nodes, viewed_by, view_of):
                     viewof_change = []
                     # Use to track view_of changes
 
-                    viewedby_add = defaultdict(lambda: [])
-                    viewedby_remove = defaultdict(lambda: [])
+                    viewedby_add = defaultdict(list)
+                    viewedby_remove = defaultdict(list)
                     # Use to track viewed_by changes
 
                     for var in node.outputs: