pytensor 2.34.0__tar.gz → 2.35.0__tar.gz

{pytensor-2.34.0/pytensor.egg-info → pytensor-2.35.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.34.0
+Version: 2.35.0
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause
@@ -21,16 +21,15 @@ Classifier: Operating System :: POSIX
 Classifier: Operating System :: Unix
 Classifier: Operating System :: MacOS
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
-Requires-Python: <3.14,>=3.10
+Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/x-rst
 License-File: LICENSE.txt
 Requires-Dist: setuptools>=59.0.0
 Requires-Dist: scipy<2,>=1
-Requires-Dist: numpy>=1.17.0
+Requires-Dist: numpy>=2.0
 Requires-Dist: filelock>=3.15
 Requires-Dist: etuples
 Requires-Dist: logical-unification

pytensor-2.35.0/doc/_drafts/benchmark_mlx_v_jax_corrected.ipynb

@@ -0,0 +1,436 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {},
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Obtaining file:///Users/carlostrujillo/Documents/GitHub/pytensor\n",
+            "  Installing build dependencies ... \u001b[?25ldone\n",
+            "\u001b[?25h  Checking if build backend supports build_editable ... \u001b[?25ldone\n",
+            "\u001b[?25h  Getting requirements to build editable ... \u001b[?25ldone\n",
+            "\u001b[?25h  Preparing editable metadata (pyproject.toml) ... \u001b[?25ldone\n",
+            "\u001b[?25hBuilding wheels for collected packages: pytensor\n",
+            "  Building editable for pytensor (pyproject.toml) ... \u001b[?25ldone\n",
+            "\u001b[?25h  Created wheel for pytensor: filename=pytensor-2.31.7+80.g06ccf91ba.dirty-0.editable-cp312-cp312-macosx_11_0_arm64.whl size=7323 sha256=c09587a5f3141d49000666d2817c5a01436f13ff5a19aa3deda20f647660afee\n",
+            "  Stored in directory: /private/var/folders/f0/rbz8xs8s17n3k3f_ccp31bvh0000gn/T/pip-ephem-wheel-cache-i00nb67k/wheels/52/f6/4c/e6784e2203d5405c94db1d544248730e598e4397674416af05\n",
+            "Successfully built pytensor\n",
+            "Installing collected packages: pytensor\n",
+            "  Attempting uninstall: pytensor\n",
+            "    Found existing installation: pytensor 2.31.7+80.g06ccf91ba.dirty\n",
+            "    Uninstalling pytensor-2.31.7+80.g06ccf91ba.dirty:\n",
+            "      Successfully uninstalled pytensor-2.31.7+80.g06ccf91ba.dirty\n",
+            "Successfully installed pytensor-2.31.7+80.g06ccf91ba.dirty\n",
+            "Note: you may need to restart the kernel to use updated packages.\n"
+          ]
+        }
+      ],
+      "source": [
+        "%pip install -e ../.. --no-deps"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import time\n",
+        "import numpy as np\n",
+        "import jax\n",
+        "import jax.numpy as jnp\n",
+        "\n",
+        "import pytensor\n",
+        "import pytensor.tensor as pt\n",
+        "from pytensor.compile.function import function\n",
+        "from pytensor.compile.mode import Mode\n",
+        "from pytensor.graph import RewriteDatabaseQuery\n",
+        "from pytensor.link.jax import JAXLinker\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Configure JAX to use float32 for consistency with MLX\n",
+        "jax.config.update(\"jax_enable_x64\", False)\n",
+        "\n",
+        "# Set up PyTensor JAX mode\n",
+        "jax_optimizer = RewriteDatabaseQuery(include=[\"jax\"], exclude=[])\n",
+        "pytensor_jax_mode = \"JAX\"\n",
+        "\n",
+        "# Try to set up MLX mode\n",
+        "try:\n",
+        "    from pytensor.link.mlx import MLXLinker\n",
+        "    import mlx.core as mx\n",
+        "    mlx_optimizer = RewriteDatabaseQuery(include=[\"mlx\"], exclude=[])\n",
+        "    pytensor_mlx_mode = \"MLX\"\n",
+        "    MLX_AVAILABLE = True\n",
+        "except ImportError:\n",
+        "    MLX_AVAILABLE = False\n",
+        "\n",
+        "def timer_jax(func, N=1000):\n",
+        "    \"\"\"Time function execution with proper JAX synchronization, repeated N times\"\"\"\n",
+        "    def wrapper(*args, **kwargs):\n",
+        "        times = []\n",
+        "        for _ in range(N):\n",
+        "            start = time.perf_counter()\n",
+        "            result = func(*args, **kwargs)\n",
+        "            if hasattr(result, 'block_until_ready'):\n",
+        "                result.block_until_ready()\n",
+        "            elif isinstance(result, (list, tuple)):\n",
+        "                for r in result:\n",
+        "                    if hasattr(r, 'block_until_ready'):\n",
+        "                        r.block_until_ready()\n",
+        "            end = time.perf_counter()\n",
+        "            times.append(end - start)\n",
+        "        \n",
+        "        mean_time = np.mean(times)\n",
+        "        std_time = np.std(times)\n",
+        "        return result, mean_time, std_time\n",
+        "    return wrapper\n",
+        "\n",
+        "def timer_mlx(func, N=1000):\n",
+        "    \"\"\"Time function execution with proper MLX synchronization, repeated N times\"\"\"\n",
+        "    def wrapper(*args, **kwargs):\n",
+        "        times = []\n",
+        "        for _ in range(N):\n",
+        "            start = time.perf_counter()\n",
+        "            result = func(*args, **kwargs)\n",
+        "            # For MLX, we need to use mx.eval() to force computation\n",
+        "            if MLX_AVAILABLE:\n",
+        "                if isinstance(result, (list, tuple)):\n",
+        "                    mx.eval(*result)\n",
+        "                else:\n",
+        "                    mx.eval(result)\n",
+        "            end = time.perf_counter()\n",
+        "            times.append(end - start)\n",
+        "        \n",
+        "        mean_time = np.mean(times)\n",
+        "        std_time = np.std(times)\n",
+        "        return result, mean_time, std_time\n",
+        "    return wrapper\n",
+        "\n",
+        "def run_benchmark(N=1000):\n",
+        "    \"\"\"Run comprehensive benchmark comparing PyTensor JAX vs MLX backends\"\"\"\n",
+        "    import pandas as pd\n",
+        "    \n",
+        "    sizes = [2, 4, 1080, 2080, 3080]\n",
+        "    results = []\n",
+        "    \n",
+        "    print(f\"Running benchmarks with N={N} repetitions per test...\")\n",
+        "    \n",
+        "    for size in sizes:\n",
+        "        print(f\"Testing {size}x{size} matrices...\")\n",
+        "        \n",
+        "        # Generate test matrices with fixed seed for reproducibility\n",
+        "        np.random.seed(42)\n",
+        "        A = np.random.randn(size, size).astype(np.float32)\n",
+        "        B = np.random.randn(size, size).astype(np.float32)\n",
+        "        C = np.random.randn(size, size).astype(np.float32)\n",
+        "\n",
+        "        pt_A = pt.matrix('A', dtype='float32')\n",
+        "        pt_B = pt.matrix('B', dtype='float32')  \n",
+        "        pt_C = pt.matrix('C', dtype='float32')\n",
+        "        result = pt.dot(pt.dot(pt_A, pt_B), pt_C)\n",
+        "\n",
+        "\n",
+        "        f_jax = function([pt_A, pt_B, pt_C], result, mode=pytensor_jax_mode, trust_input=True)\n",
+        "        f_mlx = function([pt_A, pt_B, pt_C], result, mode=pytensor_mlx_mode, trust_input=True)\n",
+        "        f_jax(A, B, C)\n",
+        "        f_mlx(A, B, C)\n",
+        "        \n",
+        "        # === TEST 1: Matrix Multiplication Chain ===\n",
+        "        # PyTensor + JAX backend\n",
+        "        @timer_jax\n",
+        "        def pytensor_jax_matmul():\n",
+        "            return f_jax(A, B, C)\n",
+        "        \n",
+        "        # PyTensor + MLX backend\n",
+        "        @timer_mlx\n",
+        "        def pytensor_mlx_matmul():\n",
+        "            if not MLX_AVAILABLE:\n",
+        "                return None, float('inf'), 0\n",
+        "            return f_mlx(A, B, C)\n",
+        "        \n",
+        "        # Run matrix multiplication test\n",
+        "        _, jax_mean, jax_std = pytensor_jax_matmul()\n",
+        "        try:\n",
+        "            _, mlx_mean, mlx_std = pytensor_mlx_matmul()\n",
+        "        except Exception as e:\n",
+        "            print(f\"MLX matmul error: {e}\")\n",
+        "            mlx_mean, mlx_std = float('inf'), 0\n",
+        "        \n",
+        "        # Calculate percentage improvement (positive = MLX is faster, negative = MLX is slower)\n",
+        "        if mlx_mean != float('inf') and mlx_mean > 0:\n",
+        "            speedup_percentage = ((jax_mean - mlx_mean) / jax_mean) * 100\n",
+        "            speedup_str = f'{speedup_percentage:+.1f}%'\n",
+        "        else:\n",
+        "            speedup_str = 'N/A'\n",
+        "        \n",
+        "        results.append({\n",
+        "            'Size': f'{size}x{size}',\n",
+        "            'Operation': 'Matrix Chain (A @ B @ C)',\n",
+        "            'PyTensor+JAX Mean (s)': f'{jax_mean:.6f}',\n",
+        "            'PyTensor+JAX Std (s)': f'{jax_std:.6f}',\n",
+        "            'PyTensor+MLX Mean (s)': f'{mlx_mean:.6f}' if mlx_mean != float('inf') else 'Error',\n",
+        "            'PyTensor+MLX Std (s)': f'{mlx_std:.6f}' if mlx_mean != float('inf') else 'N/A',\n",
+        "            'MLX Performance': speedup_str\n",
+        "        })\n",
+        "        \n",
+        "        # === TEST 2: Element-wise Operations ===\n",
+        "        # PyTensor + JAX\n",
+        "        result = pt.sin(pt_A) + pt.cos(pt_B)\n",
+        "        f_jax = function([pt_A, pt_B], result, mode=pytensor_jax_mode, trust_input=True)\n",
+        "        f_mlx = function([pt_A, pt_B], result, mode=pytensor_mlx_mode, trust_input=True)\n",
+        "        f_jax(A, B)\n",
+        "        f_mlx(A, B)\n",
+        "\n",
+        "        @timer_jax\n",
+        "        def pytensor_jax_elemwise():\n",
+        "            return f_jax(A, B)\n",
+        "        \n",
+        "        # PyTensor + MLX\n",
+        "        @timer_mlx\n",
+        "        def pytensor_mlx_elemwise():\n",
+        "            if not MLX_AVAILABLE:\n",
+        "                return None, float('inf'), 0\n",
+        "            return f_mlx(A, B)\n",
+        "        \n",
+        "        # Run element-wise test\n",
+        "        _, jax_mean, jax_std = pytensor_jax_elemwise()\n",
+        "        try:\n",
+        "            _, mlx_mean, mlx_std = pytensor_mlx_elemwise()\n",
+        "        except Exception as e:\n",
+        "            print(f\"MLX elemwise error: {e}\")\n",
+        "            mlx_mean, mlx_std = float('inf'), 0\n",
+        "        \n",
+        "        # Calculate percentage improvement\n",
+        "        if mlx_mean != float('inf') and mlx_mean > 0:\n",
+        "            speedup_percentage = ((jax_mean - mlx_mean) / jax_mean) * 100\n",
+        "            speedup_str = f'{speedup_percentage:+.1f}%'\n",
+        "        else:\n",
+        "            speedup_str = 'N/A'\n",
+        "        \n",
+        "        results.append({\n",
+        "            'Size': f'{size}x{size}',\n",
+        "            'Operation': 'Element-wise (sin(A) + cos(B))',\n",
+        "            'PyTensor+JAX Mean (s)': f'{jax_mean:.6f}',\n",
+        "            'PyTensor+JAX Std (s)': f'{jax_std:.6f}',\n",
+        "            'PyTensor+MLX Mean (s)': f'{mlx_mean:.6f}' if mlx_mean != float('inf') else 'Error',\n",
+        "            'PyTensor+MLX Std (s)': f'{mlx_std:.6f}' if mlx_mean != float('inf') else 'N/A',\n",
+        "            'MLX Performance': speedup_str\n",
+        "        })\n",
+        "        \n",
+        "        # === TEST 3: Matrix Addition with Broadcasting ===\n",
+        "        # PyTensor + JAX\n",
+        "        result = pt_A + pt_B.T\n",
+        "        f_jax = function([pt_A, pt_B], result, mode=pytensor_jax_mode, trust_input=True)\n",
+        "        f_mlx = function([pt_A, pt_B], result, mode=pytensor_mlx_mode, trust_input=True)\n",
+        "        f_jax(A, B)\n",
+        "        f_mlx(A, B)\n",
+        "        @timer_jax\n",
+        "        def pytensor_jax_broadcast():\n",
+        "            return f_jax(A, B)\n",
+        "        \n",
+        "        # PyTensor + MLX\n",
+        "        @timer_mlx\n",
+        "        def pytensor_mlx_broadcast():\n",
+        "            if not MLX_AVAILABLE:\n",
+        "                return None, float('inf'), 0\n",
+        "            return f_mlx(A, B)\n",
+        "        \n",
+        "        # Run broadcasting test\n",
+        "        _, jax_mean, jax_std = pytensor_jax_broadcast()\n",
+        "        try:\n",
+        "            _, mlx_mean, mlx_std = pytensor_mlx_broadcast()\n",
+        "        except Exception as e:\n",
+        "            print(f\"MLX broadcast error: {e}\")\n",
+        "            mlx_mean, mlx_std = float('inf'), 0\n",
+        "        \n",
+        "        # Calculate percentage improvement\n",
+        "        if mlx_mean != float('inf') and mlx_mean > 0:\n",
+        "            speedup_percentage = ((jax_mean - mlx_mean) / jax_mean) * 100\n",
+        "            speedup_str = f'{speedup_percentage:+.1f}%'\n",
+        "        else:\n",
+        "            speedup_str = 'N/A'\n",
+        "        \n",
+        "        results.append({\n",
+        "            'Size': f'{size}x{size}',\n",
+        "            'Operation': 'Broadcasting (A + B.T)',\n",
+        "            'PyTensor+JAX Mean (s)': f'{jax_mean:.6f}',\n",
+        "            'PyTensor+JAX Std (s)': f'{jax_std:.6f}',\n",
+        "            'PyTensor+MLX Mean (s)': f'{mlx_mean:.6f}' if mlx_mean != float('inf') else 'Error',\n",
+        "            'PyTensor+MLX Std (s)': f'{mlx_std:.6f}' if mlx_mean != float('inf') else 'N/A',\n",
+        "            'MLX Performance': speedup_str\n",
+        "        })\n",
+        "    \n",
+        "    # Create and display results table\n",
+        "    df = pd.DataFrame(results)\n",
+        "    return df\n",
+        "\n",
+        "def main(N=1000):\n",
+        "    \"\"\"Main benchmark execution\"\"\"\n",
+        "    # Display system info\n",
+        "    system_info = {\n",
+        "        'JAX version': jax.__version__,\n",
+        "        'PyTensor version': pytensor.__version__,\n",
+        "        'MLX Available': 'Yes' if MLX_AVAILABLE else 'No',\n",
+        "        'Platform': 'Apple Silicon' if MLX_AVAILABLE else 'Generic',\n",
+        "        'Repetitions (N)': N\n",
+        "    }\n",
+        "    \n",
+        "    if MLX_AVAILABLE:\n",
+        "        system_info['MLX version'] = mx.__version__\n",
+        "    \n",
+        "    import pandas as pd\n",
+        "    info_df = pd.DataFrame([system_info])\n",
+        "    \n",
+        "    # Then run benchmarks\n",
+        "    results_df = run_benchmark(N=N)\n",
+        "    \n",
+        "    return info_df, results_df\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {},
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Running benchmarks with N=150 repetitions per test...\n",
+            "Testing 2x2 matrices...\n",
+            "Testing 4x4 matrices...\n",
+            "Testing 1080x1080 matrices...\n",
+            "Testing 2080x2080 matrices...\n",
+            "Testing 3080x3080 matrices...\n"
+          ]
+        }
+      ],
+      "source": [
+        "iteration=150\n",
+        "_, results = main(N=iteration)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 4,
+      "metadata": {},
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "Benchmark Results over 150 repetitions:\n",
+            "     Size                      Operation PyTensor+JAX Mean (s) PyTensor+JAX Std (s) PyTensor+MLX Mean (s) PyTensor+MLX Std (s) MLX Performance\n",
+            "      2x2       Matrix Chain (A @ B @ C)              0.000009             0.000002              0.000305             0.000299        -3213.5%\n",
+            "      2x2 Element-wise (sin(A) + cos(B))              0.000007             0.000002              0.000352             0.003757        -5078.0%\n",
+            "      2x2         Broadcasting (A + B.T)              0.000007             0.000001              0.000188             0.000153        -2721.1%\n",
+            "      4x4       Matrix Chain (A @ B @ C)              0.000009             0.000001              0.000209             0.000063        -2126.2%\n",
+            "      4x4 Element-wise (sin(A) + cos(B))              0.000007             0.000001              0.000180             0.000066        -2449.5%\n",
+            "      4x4         Broadcasting (A + B.T)              0.000007             0.000003              0.000181             0.000065        -2564.1%\n",
+            "1080x1080       Matrix Chain (A @ B @ C)              0.005951             0.000356              0.001355             0.000392          +77.2%\n",
+            "1080x1080 Element-wise (sin(A) + cos(B))              0.002820             0.000107              0.000432             0.000207          +84.7%\n",
+            "1080x1080         Broadcasting (A + B.T)              0.000212             0.000035              0.000428             0.000206         -102.0%\n",
+            "2080x2080       Matrix Chain (A @ B @ C)              0.027609             0.001255              0.004550             0.002528          +83.5%\n",
+            "2080x2080 Element-wise (sin(A) + cos(B))              0.010086             0.000417              0.001175             0.000350          +88.3%\n",
+            "2080x2080         Broadcasting (A + B.T)              0.000856             0.000068              0.001124             0.000241          -31.2%\n",
+            "3080x3080       Matrix Chain (A @ B @ C)              0.093115             0.003823              0.013649             0.000513          +85.3%\n",
+            "3080x3080 Element-wise (sin(A) + cos(B))              0.022586             0.000756              0.001930             0.000287          +91.5%\n",
+            "3080x3080         Broadcasting (A + B.T)              0.002580             0.000161              0.001937             0.000257          +24.9%\n"
+          ]
+        }
+      ],
+      "source": [
+        "print(f\"\\nBenchmark Results over {iteration} repetitions:\")\n",
+        "print(results.to_string(index=False))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# # Additional timing analysis - separate compilation vs execution time\n",
+        "# if MLX_AVAILABLE:\n",
+        "#     print(\"\\n=== Detailed MLX Timing Analysis ===\")\n",
+        "    \n",
+        "#     # Test with medium-sized matrix\n",
+        "#     np.random.seed(42)\n",
+        "#     A = np.random.randn(512, 512).astype(np.float32)\n",
+        "#     B = np.random.randn(512, 512).astype(np.float32)\n",
+        "#     C = np.random.randn(512, 512).astype(np.float32)\n",
+        "    \n",
+        "#     # Create PyTensor function (compilation time)\n",
+        "#     start = time.perf_counter()\n",
+        "#     pt_A = pt.matrix('A', dtype='float32')\n",
+        "#     pt_B = pt.matrix('B', dtype='float32')\n",
+        "#     pt_C = pt.matrix('C', dtype='float32')\n",
+        "#     result_expr = pt_A @ pt_B @ pt_C\n",
+        "#     f_mlx = function([pt_A, pt_B, pt_C], result_expr, mode=pytensor_mlx_mode)\n",
+        "#     compilation_time = time.perf_counter() - start\n",
+        "    \n",
+        "#     # First execution (may include additional compilation/optimization)\n",
+        "#     start = time.perf_counter()\n",
+        "#     result = f_mlx(A, B, C)\n",
+        "#     mx.eval(result)  # Force evaluation\n",
+        "#     first_exec_time = time.perf_counter() - start\n",
+        "    \n",
+        "#     # Subsequent executions (should be faster)\n",
+        "#     exec_times = []\n",
+        "#     for _ in range(1000):\n",
+        "#         start = time.perf_counter()\n",
+        "#         result = f_mlx(A, B, C)\n",
+        "#         mx.eval(result)\n",
+        "#         exec_times.append(time.perf_counter() - start)\n",
+        "    \n",
+        "#     avg_exec_time = np.mean(exec_times)\n",
+        "#     std_exec_time = np.std(exec_times)\n",
+        "    \n",
+        "#     print(f\"Compilation time: {compilation_time:.4f}s\")\n",
+        "#     print(f\"First execution: {first_exec_time:.4f}s\")\n",
+        "#     print(f\"Average execution (5 runs): {avg_exec_time:.4f}s ± {std_exec_time:.4f}s\")\n",
+        "#     print(f\"Individual execution times: {[f'{t:.4f}' for t in exec_times]}\")\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": []
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "mlx_env",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.12.2"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 2
+}

{pytensor-2.34.0 → pytensor-2.35.0}/doc/environment.yml

@@ -3,7 +3,7 @@ channels:
   - conda-forge
   - nodefaults
 dependencies:
-  - python=3.10
+  - python=3.11
   - gcc_linux-64
   - gxx_linux-64
   - numpy

{pytensor-2.34.0 → pytensor-2.35.0}/pyproject.toml

@@ -2,7 +2,7 @@
 requires = [
     "setuptools>=59.0.0",
     "cython",
-    "numpy>=1.17.0",
+    "numpy>=2.0",
     "versioneer[toml]==0.29",
 ]
 build-backend = "setuptools.build_meta"
@@ -10,7 +10,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "pytensor"
 dynamic = ['version']
-requires-python = ">=3.10,<3.14"
+requires-python = ">=3.11,<3.14"
 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.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
@@ -49,7 +48,7 @@ keywords = [
 dependencies = [
     "setuptools>=59.0.0",
     "scipy>=1,<2",
-    "numpy>=1.17.0",
+    "numpy>=2.0",
     "filelock>=3.15",
     "etuples",
     "logical-unification",
@@ -169,7 +168,7 @@ lines-after-imports = 2
 
 
 [tool.mypy]
-python_version = "3.10"
+python_version = "3.11"
 ignore_missing_imports = true
 strict_equality = true
 warn_redundant_casts = true

{pytensor-2.34.0 → pytensor-2.35.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-10-07T21:55:32+0200",
+ "date": "2025-10-14T09:22:05+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "b6a9b68e9d9bc14e0b5d095bcbd2756895fe221d",
- "version": "2.34.0"
+ "full-revisionid": "f772066a9691c64430e100581ab1c2398e43451e",
+ "version": "2.35.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.34.0 → pytensor-2.35.0}/pytensor/bin/pytensor_cache.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import sys
+from pathlib import Path
 
 
 if sys.platform == "win32":
@@ -24,7 +25,7 @@ _logger = logging.getLogger("pytensor.bin.pytensor-cache")
 
 def print_help(exit_status):
     if exit_status:
-        print(f"command \"{' '.join(sys.argv)}\" not recognized")
+        print(f'command "{" ".join(sys.argv)}" not recognized')
     print('Type "pytensor-cache" to print the cache location')
     print('Type "pytensor-cache help" to print this help')
     print('Type "pytensor-cache clear" to erase the cache')
@@ -65,11 +66,7 @@ def main():
             # Print a warning if some cached modules were not removed, so that the
             # user knows he should manually delete them, or call
             # pytensor-cache purge, # to properly clear the cache.
-            items = [
-                item
-                for item in sorted(os.listdir(cache.dirname))
-                if item.startswith("tmp")
-            ]
+            items = list(Path(cache.dirname).glob("tmp*"))
             if items:
                 _logger.warning(
                     "There remain elements in the cache dir that you may "

{pytensor-2.34.0 → pytensor-2.35.0}/pytensor/compile/builders.py

@@ -122,7 +122,7 @@ def construct_nominal_fgraph(
     (
         local_inputs,
         local_outputs,
-        (clone_d, update_d, update_expr, new_shared_inputs),
+        (_clone_d, update_d, update_expr, new_shared_inputs),
     ) = new
 
     assert len(local_inputs) == len(inputs) + len(implicit_shared_inputs)

{pytensor-2.34.0 → pytensor-2.35.0}/pytensor/compile/function/__init__.py

@@ -12,7 +12,7 @@ from pytensor.compile.profiling import ProfileStats
 from pytensor.graph import Variable
 
 
-__all__ = ["types", "pfunc"]
+__all__ = ["pfunc", "types"]
 
 __docformat__ = "restructuredtext en"
 _logger = logging.getLogger("pytensor.compile.function")

{pytensor-2.34.0 → pytensor-2.35.0}/pytensor/compile/function/pfunc.py

@@ -328,8 +328,7 @@ def rebuild_collect_shared(
             cloned_outputs = []  # TODO: get Function.__call__ to return None
         else:
             raise TypeError(
-                "output must be an PyTensor Variable or Out "
-                "instance (or list of them)",
+                "output must be an PyTensor Variable or Out instance (or list of them)",
                 outputs,
             )
 
@@ -592,7 +591,7 @@ def construct_pfunc_ins_and_outs(
             clone_inner_graphs=True,
         )
         input_variables, cloned_extended_outputs, other_stuff = output_vars
-        clone_d, update_d, update_expr, shared_inputs = other_stuff
+        clone_d, update_d, _update_expr, shared_inputs = other_stuff
 
         # Recover only the clones of the original outputs
         if outputs is None:

{pytensor-2.34.0 → pytensor-2.35.0}/pytensor/compile/function/types.py

@@ -215,7 +215,7 @@ def add_supervisor_to_fgraph(
             input
             for spec, input in zip(input_specs, fgraph.inputs, strict=True)
             if not (
-                spec.mutable or has_destroy_handler and fgraph.has_destroyers([input])
+                spec.mutable or (has_destroy_handler and fgraph.has_destroyers([input]))
             )
         )
     )

{pytensor-2.34.0 → pytensor-2.35.0}/pytensor/compile/mode.py

@@ -27,6 +27,7 @@ from pytensor.graph.rewriting.db import (
 from pytensor.link.basic import Linker, PerformLinker
 from pytensor.link.c.basic import CLinker, OpWiseCLinker
 from pytensor.link.jax.linker import JAXLinker
+from pytensor.link.mlx.linker import MLXLinker
 from pytensor.link.numba.linker import NumbaLinker
 from pytensor.link.pytorch.linker import PytorchLinker
 from pytensor.link.vm import VMLinker
@@ -50,6 +51,7 @@ predefined_linkers = {
     "jax": JAXLinker(),
     "pytorch": PytorchLinker(),
     "numba": NumbaLinker(),
+    "mlx": MLXLinker(),
 }
 
 
@@ -407,7 +409,7 @@ class Mode:
             optimizations.
         """
 
-        link, opt = self.get_linker_optimizer(
+        _link, opt = self.get_linker_optimizer(
             self.provided_linker, self.provided_optimizer
         )
         return self.clone(optimizer=opt.register(*optimizations))
@@ -504,6 +506,20 @@ PYTORCH = Mode(
     ),
 )
 
+MLX = Mode(
+    MLXLinker(),
+    RewriteDatabaseQuery(
+        include=["fast_run"],
+        exclude=[
+            "cxx_only",
+            "BlasOpt",
+            "fusion",
+            "inplace",
+            "scan_save_mem_prealloc",
+        ],
+    ),
+)
+
 
 predefined_modes = {
     "FAST_COMPILE": FAST_COMPILE,
@@ -511,6 +527,7 @@ predefined_modes = {
     "JAX": JAX,
     "NUMBA": NUMBA,
     "PYTORCH": PYTORCH,
+    "MLX": MLX,
 }
 
 _CACHED_RUNTIME_MODES: dict[str, Mode] = {}