Add a Numba-enabled JupyterLite deployment

.github/jupyterlite/pytensor_numba.ipynb

@@ -0,0 +1,109 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "7fb27b941602401d91542211134fc71a",
+   "metadata": {},
+   "source": [
+    "# PyTensor and Numba in JupyterLite\n",
+    "\n",
+    "This example is adapted from PyTensor's introduction notebook and runs with the Numba linker entirely in the browser."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "acae54e37e7d407bbb7b55eff062a284",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "from importlib.metadata import version\n",
+    "\n",
+    "import numba\n",
+    "import numpy as np\n",
+    "\n",
+    "import pytensor\n",
+    "import pytensor.tensor as pt\n",
+    "\n",
+    "\n",
+    "assert sys.platform == \"emscripten\"\n",
+    "{\n",
+    "    \"Platform\": sys.platform,\n",
+    "    \"PyTensor\": version(\"pytensor\"),\n",
+    "    \"Numba\": numba.__version__,\n",
+    "    \"NumPy\": np.__version__,\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9a63283cbaf04dbcab1f6479b197f3a8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x = pt.vector(\"x\", shape=(None,))\n",
+    "z = pt.exp(pt.sin(x))\n",
+    "out = pt.cos((z[None, :] @ z[:, None]).squeeze())\n",
+    "\n",
+    "out.dprint()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8dd0d8092fe74a7c96281538738b07e2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "numba_fn = pytensor.function([x], out, mode=\"NUMBA\")\n",
+    "numba_fn.dprint(print_destroy_map=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "72eea5119410473aa328ad9291626812",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "values = np.array([0.25, -0.5, 1.0])\n",
+    "result = numba_fn(values)\n",
+    "expected_z = np.exp(np.sin(values))\n",
+    "expected = np.cos(expected_z @ expected_z)\n",
+    "\n",
+    "np.testing.assert_allclose(result, expected)\n",
+    "{\"Result\": result, \"Expected\": expected}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8edb47106e1a46a883d545849b8ab81b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Reuse the compiled function with a different vector length.\n",
+    "values = np.linspace(-1.0, 1.0, 5)\n",
+    "result = numba_fn(values)\n",
+    "expected_z = np.exp(np.sin(values))\n",
+    "np.testing.assert_allclose(result, np.cos(expected_z @ expected_z))\n",
+    "result"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python (XPython)",
+   "language": "python",
+   "name": "xpython"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

.github/workflows/deploy-jupyterlite.yml

@@ -0,0 +1,84 @@
+name: Build and deploy JupyterLite
+
+on:
+  workflow_dispatch:
+  pull_request:
+  push:
+    branches:
+      - main
+
+permissions: {}
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+
+    steps:
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          fetch-depth: 0
+          persist-credentials: false
+
+      - name: Install build environment
+        uses: mamba-org/setup-micromamba@d7c9bd84e824b79d2af72a2d4196c7f4300d3476 # v3.0.0
+        with:
+          environment-file: environment-wasm-build.yml
+          environment-name: pytensor-wasm-build
+          init-shell: bash
+
+      - name: Build PyTensor from this checkout
+        shell: bash -l {0}
+        run: |
+          set -euxo pipefail
+          PYTENSOR_PURE_PYTHON=1 python -m build --wheel --no-isolation
+
+          mkdir -p build/pytensor-wheel
+          wheels=(dist/pytensor-*.whl)
+          test "${#wheels[@]}" -eq 1
+          python -m zipfile -e "${wheels[0]}" build/pytensor-wheel
+
+      - name: Create the WebAssembly environment
+        shell: bash -l {0}
+        run: |
+          set -euxo pipefail
+          micromamba create -y \
+            -f environment-wasm-host.yml \
+            --platform=emscripten-wasm32
+
+          echo "PYTENSOR_WASM_PREFIX=${MAMBA_ROOT_PREFIX}/envs/pytensor-wasm-host" >> "${GITHUB_ENV}"
+
+      - name: Build JupyterLite
+        shell: bash -l {0}
+        run: |
+          set -euxo pipefail
+          jupyter lite build \
+            --XeusAddon.prefix="${PYTENSOR_WASM_PREFIX}" \
+            --XeusAddon.default_channels=https://repo.prefix.dev/emscripten-forge-4x \
+            --XeusAddon.default_channels=https://repo.prefix.dev/conda-forge \
+            --XeusAddon.mounts="$(pwd)/build/pytensor-wheel:/lib/python3.13/site-packages" \
+            --contents=.github/jupyterlite/pytensor_numba.ipynb \
+            --output-dir=dist-jupyterlite
+
+      - name: Upload Pages artifact
+        uses: actions/upload-pages-artifact@7b1f4a764d45c48632c6b24a0339c27f5614fb0b # v4.0.0
+        with:
+          path: dist-jupyterlite
+
+  deploy:
+    needs: build
+    if: github.ref == 'refs/heads/main'
+    runs-on: ubuntu-latest
+    permissions:
+      pages: write
+      id-token: write
+
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+
+    steps:
+      - name: Deploy to GitHub Pages
+        id: deployment
+        uses: actions/deploy-pages@d6db90164ac5ed86f2b6aed7e0febac5b3c0c03e # v4.0.5

environment-wasm-build.yml

@@ -0,0 +1,15 @@
+name: pytensor-wasm-build
+
+channels:
+  - conda-forge
+
+dependencies:
+  - python=3.13
+  - python-build
+  - setuptools
+  - cython
+  - numpy=2.4.*
+  - versioneer=0.29
+  - jupyterlite-core
+  - jupyterlite-xeus
+  - jupyter_server

environment-wasm-host.yml

@@ -0,0 +1,19 @@
+name: pytensor-wasm-host
+
+channels:
+  - https://prefix.dev/emscripten-forge-4x
+  - https://prefix.dev/conda-forge
+
+dependencies:
+  - python=3.13
+  - xeus-python
+  - numpy=2.4.*
+  - scipy
+  - numba=0.66.*
+  - llvmlite=0.48.*
+  - setuptools
+  - filelock
+  - etuples
+  - logical-unification
+  - minikanren
+  - cons

pytensor/link/numba/dispatch/vectorize_codegen.py

@@ -307,7 +307,7 @@ def compute_itershape(
     broadcast_pattern: tuple[tuple[bool, ...], ...],
     size: list[ir.Instruction] | None,
 ):
-    one = ir.IntType(64)(1)
+    one = ctx.get_constant(types.intp, 1)
     batch_ndim = len(broadcast_pattern[0])
     shape = [None] * batch_ndim
     if size is not None:
@@ -417,7 +417,7 @@ def make_outputs(
     """
     output_arrays = []
     output_arry_types = []
-    one = ir.IntType(64)(1)
+    one = ctx.get_constant(types.intp, 1)
     inplace_dict = dict(inplace)
     for i, (core_shape, bc, dtype) in enumerate(
         zip(output_core_shapes, out_bc, dtypes, strict=True)
@@ -448,7 +448,7 @@ def make_outputs(
             # reduction identity.  A flat scan over every element is valid
             # regardless of which axes are reduced, and seeds each kept-axis
             # accumulator cell (size-1 reduced axes included).
-            nitems = ir.IntType(64)(1)
+            nitems = ctx.get_constant(types.intp, 1)
             for dim_len in shape:
                 nitems = builder.mul(nitems, dim_len)
             ident = ctx.get_constant(dtype, reduce_identities[i])
@@ -522,7 +522,7 @@ def make_loop_call(
     ]
     destroyed_inputs = {in_idx: out_idx for out_idx, in_idx in inplace}
 
-    zero = ir.Constant(ir.IntType(64), 0)
+    zero = context.get_constant(types.intp, 0)
 
     def _wrap_negative_index(idx_val, dim_size, signed):
         """Wrap a negative index by adding the dimension size: idx + size if idx < 0.
@@ -580,12 +580,7 @@ def _wrap_negative_index(idx_val, dim_size, signed):
             val = builder.load(ptr)
             val.set_metadata("alias.scope", input_scope_set)
             val.set_metadata("noalias", output_scope_set)
-            i64 = ir.IntType(64)
-            if val.type != i64:
-                if idx_arr_type.dtype.signed:
-                    val = builder.sext(val, i64)
-                else:
-                    val = builder.zext(val, i64)
+            val = context.cast(builder, val, idx_arr_type.dtype, types.intp)
             indirect_idxs.append(val)
 
     # Load values from input arrays
@@ -1076,7 +1071,7 @@ def codegen(ctx, builder, sig, args):
             cgutils.unpack_tuple(builder, idx_arrs[k].shape) for k in range(n_indices)
         ]
 
-        one = ir.IntType(64)(1)
+        one = ctx.get_constant(types.intp, 1)
         iter_shapes = list(in_shapes)
         iter_bc = list(input_bc_patterns)
 

setup.py

@@ -13,12 +13,14 @@
 
 NAME: str = dist.get_name()  # type: ignore
 
-# Check if building for Pyodide
-is_pyodide = os.getenv("PYODIDE", "0") == "1"
-
-if is_pyodide:
-    # For pyodide we build a universal wheel that must be pure-python
-    # so we must omit the cython-version of scan.
+# Build without optional compiled extensions. Keep PYODIDE as a compatibility
+# alias for existing downstream builds.
+is_pure_python = (
+    os.getenv("PYTENSOR_PURE_PYTHON", "0") == "1" or os.getenv("PYODIDE", "0") == "1"
+)
+
+if is_pure_python:
+    # Omit the optional Cython implementation of scan.
     ext_modules = []
 else:
     ext_modules = [

tests/link/numba/test_vectorize_codegen.py

@@ -0,0 +1,40 @@
+from llvmlite import ir
+from numba import types
+
+from pytensor.link.numba.dispatch.vectorize_codegen import compute_itershape
+
+
+class Mock32BitContext:
+    """Minimal Numba context whose ``intp`` LLVM representation is i32."""
+
+    class CallConv:
+        @staticmethod
+        def return_user_exc(builder, exc, args):
+            pass
+
+    call_conv = CallConv()
+
+    @staticmethod
+    def get_constant(typ, value):
+        assert typ is types.intp
+        return ir.Constant(ir.IntType(32), value)
+
+
+def test_compute_itershape_uses_target_intp_width():
+    module = ir.Module()
+    function_type = ir.FunctionType(ir.VoidType(), ())
+    function = ir.Function(module, function_type, "compute_itershape")
+    builder = ir.IRBuilder(function.append_basic_block("entry"))
+
+    one_i32 = ir.Constant(ir.IntType(32), 1)
+    shape = compute_itershape(
+        Mock32BitContext(),
+        builder,
+        in_shapes=[[one_i32]],
+        broadcast_pattern=((True,),),
+        size=None,
+    )
+    builder.ret_void()
+
+    assert shape == [one_i32]
+    assert "icmp ne i32" in str(module)