diff --git a/.github/jupyterlite/pytensor_numba.ipynb b/.github/jupyterlite/pytensor_numba.ipynb new file mode 100644 index 0000000000..a910c9b5a0 --- /dev/null +++ b/.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 +} diff --git a/.github/workflows/deploy-jupyterlite.yml b/.github/workflows/deploy-jupyterlite.yml new file mode 100644 index 0000000000..edb7442393 --- /dev/null +++ b/.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 diff --git a/environment-wasm-build.yml b/environment-wasm-build.yml new file mode 100644 index 0000000000..c33bfc1442 --- /dev/null +++ b/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 diff --git a/environment-wasm-host.yml b/environment-wasm-host.yml new file mode 100644 index 0000000000..2c5e6b5bb4 --- /dev/null +++ b/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 diff --git a/pytensor/link/numba/dispatch/vectorize_codegen.py b/pytensor/link/numba/dispatch/vectorize_codegen.py index 9f1eb7ea75..02e68d64b1 100644 --- a/pytensor/link/numba/dispatch/vectorize_codegen.py +++ b/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) diff --git a/setup.py b/setup.py index 09202a658c..7861619410 100755 --- a/setup.py +++ b/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 = [ diff --git a/tests/link/numba/test_vectorize_codegen.py b/tests/link/numba/test_vectorize_codegen.py new file mode 100644 index 0000000000..8161027362 --- /dev/null +++ b/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)