pytensor 2.33.0__tar.gz → 2.35.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.33.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.33.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
@@ -25,4 +25,4 @@ dependencies:
   - ablog
   - pip
   - pip:
-    - -e ..
+    - -e ..[jax]

{pytensor-2.33.0 → pytensor-2.35.0}/doc/extending/creating_an_op.rst

@@ -803,10 +803,10 @@ You can omit the :meth:`Rop` functions. Try to implement the testing apparatus d
 :download:`Solution<extending_pytensor_solution_1.py>`
 
 
-:func:`as_op`
+:func:`wrap_py`
 -------------
 
-:func:`as_op` is a Python decorator that converts a Python function into a
+:func:`wrap_py` is a Python decorator that converts a Python function into a
 basic PyTensor :class:`Op` that will call the supplied function during execution.
 
 This isn't the recommended way to build an :class:`Op`, but allows for a quick implementation.
@@ -839,11 +839,11 @@ It takes an optional :meth:`infer_shape` parameter that must have this signature
     inputs PyTensor variables that were declared.
 
 .. note::
-    The python function wrapped by the :func:`as_op` decorator needs to return a new
+    The python function wrapped by the :func:`wrap_py` decorator needs to return a new
     data allocation, no views or in place modification of the input.
 
 
-:func:`as_op` Example
+:func:`wrap_py` Example
 ^^^^^^^^^^^^^^^^^^^^^
 
 .. testcode:: asop
@@ -852,14 +852,14 @@ It takes an optional :meth:`infer_shape` parameter that must have this signature
     import pytensor.tensor as pt
     import numpy as np
     from pytensor import function
-    from pytensor.compile.ops import as_op
+    from pytensor.compile.ops import wrap_py
 
     def infer_shape_numpy_dot(fgraph, node, input_shapes):
         ashp, bshp = input_shapes
         return [ashp[:-1] + bshp[-1:]]
 
 
-    @as_op(
+    @wrap_py(
         itypes=[pt.dmatrix, pt.dmatrix],
         otypes=[pt.dmatrix],
         infer_shape=infer_shape_numpy_dot,

{pytensor-2.33.0 → pytensor-2.35.0}/doc/extending/extending_pytensor_solution_1.py

@@ -167,9 +167,9 @@ class TestSumDiffOp(utt.InferShapeTester):
 
 import numpy as np
 
-# as_op exercice
+# wrap_py exercice
 import pytensor
-from pytensor.compile.ops import as_op
+from pytensor.compile.ops import wrap_py
 
 
 def infer_shape_numpy_dot(fgraph, node, input_shapes):
@@ -177,7 +177,7 @@ def infer_shape_numpy_dot(fgraph, node, input_shapes):
     return [ashp[:-1] + bshp[-1:]]
 
 
-@as_op(
+@wrap_py(
     itypes=[pt.fmatrix, pt.fmatrix],
     otypes=[pt.fmatrix],
     infer_shape=infer_shape_numpy_dot,
@@ -192,7 +192,7 @@ def infer_shape_numpy_add_sub(fgraph, node, input_shapes):
     return [ashp[0]]
 
 
-@as_op(
+@wrap_py(
     itypes=[pt.fmatrix, pt.fmatrix],
     otypes=[pt.fmatrix],
     infer_shape=infer_shape_numpy_add_sub,
@@ -201,7 +201,7 @@ def numpy_add(a, b):
     return np.add(a, b)
 
 
-@as_op(
+@wrap_py(
     itypes=[pt.fmatrix, pt.fmatrix],
     otypes=[pt.fmatrix],
     infer_shape=infer_shape_numpy_add_sub,

{pytensor-2.33.0 → pytensor-2.35.0}/doc/library/index.rst

@@ -61,10 +61,16 @@ Convert to Variable
 
 .. autofunction:: pytensor.as_symbolic(...)
 
+Wrap JAX functions
+==================
+
+.. autofunction:: wrap_jax(...)
+
+   Alias for :func:`pytensor.link.jax.ops.wrap_jax`
+
 Debug
 =====
 
 .. autofunction:: pytensor.dprint(...)
 
    Alias for :func:`pytensor.printing.debugprint`
-

{pytensor-2.33.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.33.0 → pytensor-2.35.0}/pytensor/__init__.py

@@ -166,7 +166,7 @@ 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
-
+from pytensor.link.jax.ops import wrap_jax
 # isort: on
 
 

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

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-10-01T04:44:52+0200",
+ "date": "2025-10-14T09:22:05+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "7779b07b892118151f96d1ed98bd69d4281c687c",
- "version": "2.33.0"
+ "full-revisionid": "f772066a9691c64430e100581ab1c2398e43451e",
+ "version": "2.35.0"
 }
 '''  # END VERSION_JSON