Skip to content

Draft: use pywebnn instead of unvalidated JSON #3

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
theHamsta wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Draft: use pywebnn instead of unvalidated JSON #3

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8104484
pyproject refactor
theHamsta Oct 1, 2026
451f4c6
Refactor
theHamsta Oct 1, 2026
622856d
Make tests more relaxed
theHamsta Oct 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
136 changes: 33 additions & 103 deletions ADDING_OPS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,120 +1,50 @@
# Adding ATen Ops
# Adding ATen operations

The exporter uses `torch.export.export` which produces an ATen IR graph. Every
`call_function` node has a target like `aten.relu.default`. The dispatch table in
`webnn_op_mappings.py` maps `str(node.target)` → a method name on
`WebNNGraphGenerator`. If the target isn't in the table the node is emitted as a
comment and the `.webnn` file will fail to parse.
`torch.export.export` produces ATen `call_function` nodes. The dispatch table in
`webnn_torch_export/webnn_op_mappings.py` maps `str(node.target)` to a method on
`WebNNGraphGenerator`. Unsupported targets fail immediately.

## Quick reference: adding an op

**1. Add an entry to `ATEN_OP_TABLE` in `webnn_op_mappings.py`:**

```python
"aten.some_op.overload": "_convert_some_op",
```

**2. Implement `_convert_some_op` in `webnn_generator.py`:**
Add the target to `ATEN_OP_TABLE` and implement its converter. Converters receive
`(node, args, kwargs)` with FX references already resolved into `CheckedOperand`
objects. Scalars and nested argument structures are preserved.

```python
def _convert_some_op(self, node: fx.Node, output: str, inputs: List[str]) -> str:
x = inputs[0]
# node.args holds positional args (first is always the input tensor node)
# node.kwargs holds keyword args
return f"[{output}] = webnnOp({x});"
```

The method receives:
- `output` — pre-allocated operand name for this node's result
- `inputs` — operand names for every `fx.Node` in `node.args`, in order
- `node.args` / `node.kwargs` — raw args (scalars, lists, bools) from the ATen schema

Return `None` instead of a string to suppress the node entry entirely (use this for
ops that bake a constant into `inline_constants` and override `node_to_operand`, like
`_convert_arange`).

**3. To find the exact target string for an unfamiliar op:**

```python
ep = torch.export.export(model, (x,))
for n in ep.graph.nodes:
if n.op == "call_function":
print(str(n.target), n.args, n.kwargs)
```

---
# webnn_op_mappings.py
"aten.some_op.default": "_some_op",

## Example patterns
# webnn_generator.py

### No-op
```python
# webnn_op_mappings.py
"aten.contiguous.default": "_convert_identity",
def _some_op(self, node, args, kwargs):
x = args[0]
return self._emit("some_webnn_op", torch.some_op, x)
```

### Direct WebNN equivalent
```python
# webnn_op_mappings.py
"aten.bmm.default": "_convert_matmul",
```
`_convert_matmul` already exists; nothing else needed.
Always create intermediate operands through `_emit`, `_constant`, or existing
checked helpers. `_emit` evaluates the reference function on PyTorch meta tensors,
calls the named pywebnn builder method, and compares the resulting operand shape
immediately. Builder options are keyword arguments; capture the corresponding
options in the reference function:

### Parameterised activation (e.g. relu6)
```python
# webnn_op_mappings.py
"aten.relu6.default": "_convert_relu6",

# webnn_generator.py
def _convert_relu6(self, node: fx.Node, output: str, inputs: List[str]) -> str:
return f"[{output}] = clamp({inputs[0]}, minValue=0.0, maxValue=6.0);"
return self._emit(
"softmax", lambda t: torch.softmax(t, axis), x, axis=axis
)
```

### Scalar-first subtract (`1 - x`)
```python
# webnn_op_mappings.py
"aten.rsub.Scalar": "_convert_rsub_scalar",
The reference function receives meta tensors in place of operand arguments.
`_emit`'s keyword options are passed only to the WebNN builder. The FX result is
then checked against `node.meta["val"]` before downstream nodes are processed.
Return tuples/lists for multi-output operations; every represented output is
checked. Unsupported auxiliary outputs may be `None` only if unused; consuming
one must raise an error.

# webnn_generator.py
def _convert_rsub_scalar(self, node: fx.Node, output: str, inputs: List[str]) -> str:
# rsub(x, scalar) = scalar - x
scalar = node.args[1]
c = self._create_inline_constant(float(scalar))
return f"[{output}] = sub({c}, {inputs[0]});"
```
For a decomposition, check every step:

### Op that produces a constant (no graph node needed)
```python
# webnn_op_mappings.py
"aten.full.default": "_convert_full",

# webnn_generator.py — return None to skip nodes {} entry
def _convert_full(self, node: fx.Node, output: str, inputs: List[str]) -> Optional[str]:
size = list(node.args[0])
fill = node.args[1]
dtype = node.kwargs.get("dtype", torch.float32) or torch.float32
values = torch.full(size, fill, dtype=dtype)
name = f"const_full_{self.operand_counter}"; self.operand_counter += 1
self.inline_constants[name] = values
self.operand_shapes[name] = list(values.shape)
self.node_to_operand[node.name] = name # redirect downstream refs
return None
sigmoid = self._emit("sigmoid", torch.sigmoid, x)
return self._binary("mul", x, sigmoid)
```

### Multi-step decomposition
```python
def _convert_select_int(self, node: fx.Node, output: str, inputs: List[str]) -> str:
# aten.select.int(input, dim, index) → slice dim then squeeze
x = inputs[0]
in_shape = self._get_node_shape(node.args[0])
dim = int(node.args[1]) % len(in_shape)
idx = int(node.args[2])
starts = [0] * len(in_shape); starts[dim] = idx
sizes = list(in_shape); sizes[dim] = 1
out_shape = in_shape[:dim] + in_shape[dim+1:]
sliced = f"operand_{self.operand_counter}"; self.operand_counter += 1
s_str = ", ".join(map(str, starts))
sz_str = ", ".join(map(str, sizes))
os_str = ", ".join(map(str, out_shape))
return (f"[{sliced}] = slice({x}, starts=[{s_str}], sizes=[{sz_str}]);\n"
f"\t[{output}] = reshape({sliced}, newShape=[{os_str}]);")
```
Add an execution test comparing the built network with the original PyTorch
model across relevant shapes/options. Do not compare PyTorch only with its own
exported ATen program. Include a negative test when adding an unsupported option.
121 changes: 52 additions & 69 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,93 +1,76 @@
# WebNN PyTorch Exporter

Export PyTorch models to WebNN format using torch dynamo IR.

---

## [WARNING] EXPERIMENTAL - DO NOT USE IN PRODUCTION

This is an early-stage experimental implementation for research and exploration. Many features are incomplete, untested, or may change significantly.

---
Build and execute fixed-size WebNN networks directly with `pywebnn` from
`torch.export` ATen graphs.

## Installation

### For Development

```bash
# Clone your forked repository
git clone https://github.com/<yourusername>/webnn_torch_export.git
cd webnn_torch_export

# Install in editable mode with dev dependencies
pip install -e ".[dev]"
# Optional: Run pytest
pytest
```

### For Use

```bash
pip install webnn_torch_export
```
`pywebnn` is a required dependency. Use the updated local pywebnn bindings that
expose `MLGraphBuilder.new_uncompiled`, `rustnn_webnn_text_for_outputs`, and
`rustnn_save_webnn`. Construction and text export need no runtime backend. For
inference, the selected runtime must be available to the dynamic loader. For ONNX Runtime, put the directory containing
`libonnxruntime.so` in `LD_LIBRARY_PATH`; if multiple runtimes are installed,
`WEBNN_RUNTIME_LIBRARY` selects the exact library before importing `webnn`.

### Use the Exporter in Your Code

**Export graph only:**
## Export and execute

```python
from webnn_torch_export import export_model
import torch
import torch.nn as nn

# Create your model
model = nn.Conv2d(3, 16, kernel_size=3)
input_tensor = torch.randn(1, 3, 28, 28)

# Export with debug output
compiled_model, exporter = export_model(model, input_tensor, debug=True)
from webnn_torch_export import export_model

# Save exported graph
exporter.save_to_file('my_export.json')
model = torch.nn.Sequential(torch.nn.Linear(10, 20), torch.nn.ReLU()).eval()
x = torch.randn(2, 10)

# Access exported graphs programmatically
for graph in exporter.exported_graphs:
print(graph['nodes'])
exporter, exported_program = export_model(model, x)
executor = exporter.to_executor()
y = executor(x)
torch.testing.assert_close(y, model(x), rtol=1e-4, atol=1e-4)
```

**Export graph + weights:**
`export_model` constructs `exporter.webnn_graph` immediately with an uncompiled
pywebnn builder. Parameters, buffers,
lifted tensor constants, and generated constants become builder constants. No
text graph, external weight lookup, file parsing, or temporary files are needed
for inference. `exporter.to_executor()` creates a CPU context when needed. Pass
`context=your_webnn_context` to select a backend/device.

Every operand is checked immediately against a shape independently inferred by
PyTorch, including inputs, constants, and decomposition intermediates. Each FX
node result is also checked against `torch.export` tensor metadata. A mismatch
raises `ShapeMismatchError` with the node, operation, and both shapes; unsupported
operators raise `NotImplementedError` during construction.

Networks have fixed input shapes and dtypes. Symbolic dimensions, empty tensor
operands, input/buffer mutation, and data-dependent shapes are unsupported.
Use `model.eval()` for inference. Tuple/list/dictionary tensor inputs and outputs
keep their exported structure and order. Outputs are CPU PyTorch tensors.
Some operator options, such as training dropout, pooling indices, arbitrary
einsum equations, and bilinear `align_corners=True`, are unsupported and raise
an error if consumed or requested.

## Optional persistence

```python
from webnn_torch_export import export_model_with_weights, load_weights_from_safetensors
import torch
import torch.nn as nn

# Create and export your model
model = nn.Sequential(
nn.Conv2d(3, 16, 3),
nn.ReLU(),
nn.Linear(16, 10)
)
input_tensor = torch.randn(1, 3, 28, 28)

# Export both graph and weights
compiled_model, exporter = export_model_with_weights(
model=model,
example_input=input_tensor,
graph_path="model_graph.json",
weights_path="model_weights.safetensors",
debug=False
)

# Later: load weights into a fresh model
new_model = nn.Sequential(
nn.Conv2d(3, 16, 3),
nn.ReLU(),
nn.Linear(16, 10)
)
load_weights_from_safetensors(new_model, "model_weights.safetensors")
text = exporter.to_webnn_text() # RustNN text with inline constants
exporter.save_to_webnn("model.webnn") # RustNN text + model.safetensors
exporter.save_weights("torch_weights.safetensors") # optional PyTorch state snapshot
```

`export_model_with_weights` remains a compatibility helper for saving the RustNN
text and external constant weights; `return_executor=True` wraps the already-built
network. The RustNN saver writes `<stem>.safetensors` beside the text file and
uses `@weights` references. If a different `weights_path` is requested, the helper
also copies that sidecar to the requested path. The exporter constructs operands
through pywebnn and delegates all text formatting and sidecar writing to RustNN.

The `webnn-export program.pt2 model.webnn` command builds a network from a fixed-size
program previously saved with `torch.export.save`.

## License

Apache License (2.0) (see LICENSE file)
Apache License (2.0), see LICENSE.
Loading
Loading