Add minor changes and resolve conflict

Author: AdvH039

.github/workflows/mypy.yml

@@ -13,7 +13,7 @@ jobs:
         shell: bash -leo pipefail {0}
     steps:
       - uses: actions/checkout@v4
-      - uses: mamba-org/setup-micromamba@v1
+      - uses: mamba-org/setup-micromamba@v2
         with:
           micromamba-version: "1.5.10-0" # until https://github.com/mamba-org/setup-micromamba/issues/225 is resolved
           environment-file: environment.yml

.github/workflows/pypi.yml

@@ -62,7 +62,7 @@ jobs:
           fetch-depth: 0
 
       - name: Build wheels
-        uses: pypa/cibuildwheel@v2.21.2
+        uses: pypa/cibuildwheel@v2.22.0
 
       - uses: actions/upload-artifact@v4
         with:
@@ -78,7 +78,7 @@ jobs:
           fetch-depth: 0
 
       - name: Set up Python
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: '3.11'
 
@@ -145,7 +145,7 @@ jobs:
           name: universal_wheel
           path: dist
 
-      - uses: pypa/gh-action-pypi-publish@v1.10.3
+      - uses: pypa/gh-action-pypi-publish@v1.12.2
         with:
           user: __token__
           password: ${{ secrets.pypi_password }}

.github/workflows/test.yml

@@ -61,16 +61,17 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - uses: pre-commit/[email protected]
 
-  test_ubuntu:
-    name: "Test py${{ matrix.python-version }} : fast-compile ${{ matrix.fast-compile }} : float32 ${{ matrix.float32 }} : ${{ matrix.part }}"
+  test:
+    name: "${{ matrix.os }} test py${{ matrix.python-version }} : fast-compile ${{ matrix.fast-compile }} : float32 ${{ matrix.float32 }} : ${{ matrix.part }}"
     needs:
       - changes
       - style
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.os }}
     if: ${{ needs.changes.outputs.changes == 'true' && needs.style.result == 'success' }}
     strategy:
       fail-fast: false
       matrix:
+        os: ["ubuntu-latest"]
         python-version: ["3.10", "3.12"]
         fast-compile: [0, 1]
         float32: [0, 1]
@@ -103,36 +104,50 @@ jobs:
             fast-compile: 1
         include:
           - install-numba: 1
+            os: "ubuntu-latest"
             python-version: "3.10"
             fast-compile: 0
             float32: 0
             part: "tests/link/numba"
           - install-numba: 1
+            os: "ubuntu-latest"
             python-version: "3.12"
             fast-compile: 0
             float32: 0
             part: "tests/link/numba"
           - install-jax: 1
+            os: "ubuntu-latest"
             python-version: "3.10"
             fast-compile: 0
             float32: 0
             part: "tests/link/jax"
           - install-jax: 1
+            os: "ubuntu-latest"
             python-version: "3.12"
             fast-compile: 0
             float32: 0
             part: "tests/link/jax"
           - install-torch: 1
+            os: "ubuntu-latest"
             python-version: "3.10"
             fast-compile: 0
             float32: 0
             part: "tests/link/pytorch"
+          - os: macos-15
+            python-version: "3.12"
+            fast-compile: 0
+            float32: 0
+            install-numba: 0
+            install-jax: 0
+            install-torch: 0
+            part: "tests/tensor/test_blas.py tests/tensor/test_elemwise.py tests/tensor/test_math_scipy.py"
+
     steps:
       - uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Set up Python ${{ matrix.python-version }}
-        uses: mamba-org/setup-micromamba@v1
+        uses: mamba-org/setup-micromamba@v2
         with:
           environment-name: pytensor-test
           micromamba-version: "1.5.10-0" # until https://github.com/mamba-org/setup-micromamba/issues/225 is resolved
@@ -146,15 +161,19 @@ jobs:
           MATRIX_CONTEXT: ${{ toJson(matrix) }}
         run: |
           echo $MATRIX_CONTEXT
-          export MATRIX_ID=`echo $MATRIX_CONTEXT | md5sum | cut -c 1-32`
+          export MATRIX_ID=`echo $MATRIX_CONTEXT | sha256sum | cut -c 1-32`
           echo $MATRIX_ID
           echo "id=$MATRIX_ID" >> $GITHUB_OUTPUT
 
       - name: Install dependencies
         shell: micromamba-shell {0}
         run: |
 
-          micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock
+          if [[ $OS == "macos-15" ]]; then
+            micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" numpy scipy pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock libblas=*=*accelerate;
+          else
+            micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock;
+          fi
           if [[ $INSTALL_NUMBA == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numba>=0.57"; fi
           if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax jaxlib numpyro && pip install tensorflow-probability; fi
           if [[ $INSTALL_TORCH == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" pytorch pytorch-cuda=12.1 "mkl<=2024.0" -c pytorch -c nvidia; fi
@@ -163,12 +182,17 @@ jobs:
           pip install -e ./
           micromamba list && pip freeze
           python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'
-          python -c 'import pytensor; assert pytensor.config.blas__ldflags != "", "Blas flags are empty"'
+          if [[ $OS == "macos-15" ]]; then
+            python -c 'import pytensor; assert pytensor.config.blas__ldflags.startswith("-framework Accelerate"), "Blas flags are not set to MacOS Accelerate"';
+          else
+            python -c 'import pytensor; assert pytensor.config.blas__ldflags != "", "Blas flags are empty"';
+          fi
         env:
           PYTHON_VERSION: ${{ matrix.python-version }}
           INSTALL_NUMBA: ${{ matrix.install-numba }}
           INSTALL_JAX: ${{ matrix.install-jax }}
           INSTALL_TORCH: ${{ matrix.install-torch}}
+          OS: ${{ matrix.os}}
 
       - name: Run tests
         shell: micromamba-shell {0}
@@ -206,7 +230,7 @@ jobs:
         with:
           fetch-depth: 0
       - name: Set up Python 3.10
-        uses: mamba-org/setup-micromamba@v1
+        uses: mamba-org/setup-micromamba@v2
         with:
           environment-name: pytensor-test
           micromamba-version: "1.5.10-0" # until https://github.com/mamba-org/setup-micromamba/issues/225 is resolved
@@ -249,10 +273,10 @@ jobs:
     if: ${{ always() }}
     runs-on: ubuntu-latest
     name: "All tests"
-    needs: [changes, style, test_ubuntu]
+    needs: [changes, style, test]
     steps:
       - name: Check build matrix status
-        if: ${{ needs.changes.outputs.changes == 'true' && (needs.style.result != 'success' || needs.test_ubuntu.result != 'success') }}
+        if: ${{ needs.changes.outputs.changes == 'true' && (needs.style.result != 'success' || needs.test.result != 'success') }}
         run: exit 1
 
   upload-coverage:
@@ -280,7 +304,7 @@ jobs:
           merge-multiple: true
 
       - name: Upload coverage to Codecov
-        uses: codecov/codecov-action@v4
+        uses: codecov/codecov-action@v5
         with:
           directory: ./coverage/
           fail_ci_if_error: true

.pre-commit-config.yaml

@@ -27,7 +27,7 @@ repos:
     - id: sphinx-lint 
       args: ["."]
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.7.1
+    rev: v0.7.3
     hooks:
       - id: ruff
         args: ["--fix", "--output-format=full"]

doc/library/tensor/random/index.rst

@@ -30,6 +30,7 @@ sophisticated `Op`\s like `Scan`, which makes it a user-friendly random variable
 interface in PyTensor.
 
 For an example of how to use random numbers, see :ref:`Using Random Numbers <using_random_numbers>`.
+For a technical explanation of how PyTensor implements random variables see :ref:`prng`.
 
 
 .. class:: RandomStream()

doc/tutorial/examples.rst

@@ -357,6 +357,9 @@ hold here as well. PyTensor's random objects are defined and implemented in
 :ref:`RandomStream<libdoc_tensor_random_utils>` and, at a lower level,
 in :ref:`RandomVariable<libdoc_tensor_random_basic>`.
 
+For a more technical explanation of how PyTensor implements random variables see :ref:`prng`.
+
+
 Brief Example
 -------------
 

doc/tutorial/prng.rst

@@ -5,7 +5,9 @@ Pseudo random number generation in PyTensor
 ===========================================
 
 PyTensor has native support for `pseudo random number generation (PRNG) <https://en.wikipedia.org/wiki/Pseudorandom_number_generator>`_.
-This document describes how PRNGs are implemented in PyTensor, via the RandomVariable Operator.
+
+This document describes the details of how PRNGs are implemented in PyTensor, via the RandomVariable Operator.
+For a more applied example see :ref:`using_random_numbers`
 
 We also discuss how initial seeding and seeding updates are implemented, and some harder cases such as using RandomVariables inside Scan, or with other backends like JAX.
 

pyproject.toml

@@ -129,8 +129,12 @@ exclude = ["doc/", "pytensor/_version.py"]
 docstring-code-format = true
 
 [tool.ruff.lint]
-select = ["C", "E", "F", "I", "UP", "W", "RUF", "PERF", "PTH", "ISC"]
+select = ["B905", "C", "E", "F", "I", "UP", "W", "RUF", "PERF", "PTH", "ISC"]
 ignore = ["C408", "C901", "E501", "E741", "RUF012", "PERF203", "ISC001"]
+unfixable = [
+    # zip-strict: the auto-fix adds `strict=False` but we might want `strict=True` instead
+    "B905",
+]
 
 
 [tool.ruff.lint.isort]

pytensor/__init__.py

@@ -164,6 +164,7 @@ def get_underlying_scalar_constant(v):
 from pytensor.scan import checkpoints
 from pytensor.scan.basic import scan
 from pytensor.scan.views import foldl, foldr, map, reduce
+from pytensor.compile.builders import OpFromGraph
 
 # isort: on
 

pytensor/compile/builders.py

@@ -43,15 +43,15 @@ def infer_shape(outs, inputs, input_shapes):
     # TODO: ShapeFeature should live elsewhere
     from pytensor.tensor.rewriting.shape import ShapeFeature
 
-    for inp, inp_shp in zip(inputs, input_shapes):
+    for inp, inp_shp in zip(inputs, input_shapes, strict=True):
         if inp_shp is not None and len(inp_shp) != inp.type.ndim:
             assert len(inp_shp) == inp.type.ndim
 
     shape_feature = ShapeFeature()
     shape_feature.on_attach(FunctionGraph([], []))
 
     # Initialize shape_of with the input shapes
-    for inp, inp_shp in zip(inputs, input_shapes):
+    for inp, inp_shp in zip(inputs, input_shapes, strict=True):
         shape_feature.set_shape(inp, inp_shp)
 
     def local_traverse(out):
@@ -108,7 +108,9 @@ def construct_nominal_fgraph(
 
     replacements = dict(
         zip(
-            inputs + implicit_shared_inputs, dummy_inputs + dummy_implicit_shared_inputs
+            inputs + implicit_shared_inputs,
+            dummy_inputs + dummy_implicit_shared_inputs,
+            strict=True,
         )
     )
 
@@ -138,7 +140,7 @@ def construct_nominal_fgraph(
         NominalVariable(n, var.type) for n, var in enumerate(local_inputs)
     )
 
-    fgraph.replace_all(zip(local_inputs, nominal_local_inputs))
+    fgraph.replace_all(zip(local_inputs, nominal_local_inputs, strict=True))
 
     for i, inp in enumerate(fgraph.inputs):
         nom_inp = nominal_local_inputs[i]
@@ -562,7 +564,9 @@ def lop_overrides(inps, grads):
             # compute non-overriding downsteam grads from upstreams grads
             # it's normal some input may be disconnected, thus the 'ignore'
             wrt = [
-                lin for lin, gov in zip(inner_inputs, custom_input_grads) if gov is None
+                lin
+                for lin, gov in zip(inner_inputs, custom_input_grads, strict=True)
+                if gov is None
             ]
             default_input_grads = fn_grad(wrt=wrt) if wrt else []
             input_grads = self._combine_list_overrides(
@@ -653,7 +657,7 @@ def _build_and_cache_rop_op(self):
             f = [
                 output
                 for output, custom_output_grad in zip(
-                    inner_outputs, custom_output_grads
+                    inner_outputs, custom_output_grads, strict=True
                 )
                 if custom_output_grad is None
             ]
@@ -733,18 +737,24 @@ def make_node(self, *inputs):
 
         non_shared_inputs = [
             inp_t.filter_variable(inp)
-            for inp, inp_t in zip(non_shared_inputs, self.input_types)
+            for inp, inp_t in zip(non_shared_inputs, self.input_types, strict=True)
         ]
 
         new_shared_inputs = inputs[num_expected_inps:]
-        inner_and_input_shareds = list(zip(self.shared_inputs, new_shared_inputs))
+        inner_and_input_shareds = list(
+            zip(self.shared_inputs, new_shared_inputs, strict=True)
+        )
 
         if not all(inp_s == inn_s for inn_s, inp_s in inner_and_input_shareds):
             # The shared variables are not equal to the original shared
             # variables, so we construct a new `Op` that uses the new shared
             # variables instead.
             replace = dict(
-                zip(self.inner_inputs[num_expected_inps:], new_shared_inputs)
+                zip(
+                    self.inner_inputs[num_expected_inps:],
+                    new_shared_inputs,
+                    strict=True,
+                )
             )
 
             # If the new shared variables are inconsistent with the inner-graph,
@@ -811,7 +821,7 @@ def infer_shape(self, fgraph, node, shapes):
         # each shape call. PyTensor optimizer will clean this up later, but this
         # will make extra work for the optimizer.
 
-        repl = dict(zip(self.inner_inputs, node.inputs))
+        repl = dict(zip(self.inner_inputs, node.inputs, strict=True))
         clone_out_shapes = [s for s in out_shapes if isinstance(s, tuple)]
         cloned = clone_replace(sum(clone_out_shapes, ()), replace=repl)
         ret = []
@@ -853,5 +863,6 @@ def clone(self):
     def perform(self, node, inputs, outputs):
         variables = self.fn(*inputs)
         assert len(variables) == len(outputs)
-        for output, variable in zip(outputs, variables):
+        # strict=False because asserted above
+        for output, variable in zip(outputs, variables, strict=False):
             output[0] = variable