fluxfem 0.1.1a0__py3-none-any.whl → 0.1.4__py3-none-any.whl

fluxfem/core/weakform.py

@@ -1,82 +1,312 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, Callable
+from typing import Any, Callable, Iterator, Literal, get_args
+import inspect
+
+import numpy as np
 
 import jax.numpy as jnp
 import jax
 
 from ..physics import operators as _ops
+from .context_types import FormFieldLike, ParamsLike, SurfaceContext, UElement, VolumeContext
+
+
+OpName = Literal[
+    "lit",
+    "getattr",
+    "value",
+    "grad",
+    "pow",
+    "eye",
+    "det",
+    "inv",
+    "transpose",
+    "log",
+    "surface_normal",
+    "surface_measure",
+    "volume_measure",
+    "sym_grad",
+    "outer",
+    "add",
+    "sub",
+    "mul",
+    "matmul",
+    "matmul_std",
+    "neg",
+    "dot",
+    "sdot",
+    "ddot",
+    "inner",
+    "action",
+    "gaction",
+    "transpose_last2",
+    "einsum",
+]
 
+# Use OpName as the single source of truth for valid ops.
+_OP_NAMES: frozenset[str] = frozenset(get_args(OpName))
+
+
+_PRECEDENCE: dict[str, int] = {
+    "add": 10,
+    "sub": 10,
+    "mul": 20,
+    "matmul": 20,
+    "matmul_std": 20,
+    "inner": 20,
+    "dot": 20,
+    "sdot": 20,
+    "ddot": 20,
+    "pow": 30,
+    "neg": 40,
+    "transpose": 50,
+}
+
+
+def _pretty_render_arg(arg, prec: int | None = None) -> str:
+    if isinstance(arg, Expr):
+        return _pretty_expr(arg, prec or 0)
+    if isinstance(arg, FieldRef):
+        if arg.name is None:
+            return f"{arg.role}"
+        return f"{arg.role}:{arg.name}"
+    if isinstance(arg, ParamRef):
+        return "param"
+    return repr(arg)
+
+
+def _pretty_wrap(text: str, prec: int, parent_prec: int) -> str:
+    if prec < parent_prec:
+        return f"({text})"
+    return text
+
+
+def _pretty_expr(expr: Expr, parent_prec: int = 0) -> str:
+    op = expr.op
+    args = expr.args
 
+    if op == "lit":
+        return repr(args[0])
+    if op == "getattr":
+        base = _pretty_render_arg(args[0], _PRECEDENCE.get("transpose", 50))
+        return f"{base}.{args[1]}"
+    if op == "value":
+        return f"val({_pretty_render_arg(args[0])})"
+    if op == "grad":
+        return f"grad({_pretty_render_arg(args[0])})"
+    if op == "sym_grad":
+        return f"sym_grad({_pretty_render_arg(args[0])})"
+    if op == "neg":
+        inner = _pretty_render_arg(args[0], _PRECEDENCE["neg"])
+        return _pretty_wrap(f"-{inner}", _PRECEDENCE["neg"], parent_prec)
+    if op == "transpose":
+        inner = _pretty_render_arg(args[0], _PRECEDENCE["transpose"])
+        return _pretty_wrap(f"{inner}.T", _PRECEDENCE["transpose"], parent_prec)
+    if op == "pow":
+        base = _pretty_render_arg(args[0], _PRECEDENCE["pow"])
+        exp = _pretty_render_arg(args[1], _PRECEDENCE["pow"] + 1)
+        return _pretty_wrap(f"{base}**{exp}", _PRECEDENCE["pow"], parent_prec)
+    if op in {"add", "sub", "mul", "matmul", "dot", "sdot", "ddot"}:
+        left = _pretty_render_arg(args[0], _PRECEDENCE[op])
+        right = _pretty_render_arg(args[1], _PRECEDENCE[op] + 1)
+        symbol = {
+            "add": "+",
+            "sub": "-",
+            "mul": "*",
+            "matmul": "@",
+            "inner": "|",
+            "dot": "dot",
+            "sdot": "sdot",
+            "ddot": "ddot",
+        }[op]
+        if symbol in {"dot", "sdot", "ddot"}:
+            text = f"{symbol}({left}, {right})"
+        else:
+            text = f"{left} {symbol} {right}"
+        return _pretty_wrap(text, _PRECEDENCE[op], parent_prec)
+    if op == "inner":
+        return f"inner({_pretty_render_arg(args[0])}, {_pretty_render_arg(args[1])})"
+    if op in {"action", "gaction"}:
+        return f"{op}({_pretty_render_arg(args[0])}, {_pretty_render_arg(args[1])})"
+    if op == "matmul_std":
+        return f"matmul_std({_pretty_render_arg(args[0])}, {_pretty_render_arg(args[1])})"
+    if op == "outer":
+        return f"outer({_pretty_render_arg(args[0])}, {_pretty_render_arg(args[1])})"
+    if op in {
+        "eye",
+        "det",
+        "inv",
+        "log",
+        "surface_normal",
+        "surface_measure",
+        "volume_measure",
+        "transpose_last2",
+        "einsum",
+    }:
+        rendered = ", ".join(_pretty_render_arg(arg) for arg in args)
+        return f"{op}({rendered})"
+    rendered = ", ".join(_pretty_render_arg(arg) for arg in args)
+    return f"{op}({rendered})"
+
+
+def _as_expr(obj) -> Expr | FieldRef | ParamRef:
+    """Normalize inputs into Expr/FieldRef/ParamRef nodes."""
+    if isinstance(obj, Expr):
+        return obj
+    if isinstance(obj, FieldRef):
+        return obj
+    if isinstance(obj, ParamRef):
+        return obj
+    if isinstance(obj, (int, float, bool, str)):
+        return Expr("lit", obj)
+    if isinstance(obj, np.generic):
+        return Expr("lit", obj.item())
+    if isinstance(obj, tuple):
+        try:
+            hash(obj)
+        except TypeError as exc:
+            raise TypeError(
+                "Expr tuple literal must be hashable; use only immutable items."
+            ) from exc
+        return Expr("lit", obj)
+    raise TypeError(
+        "Expr literal must be a scalar or hashable tuple. "
+        "Arrays are not allowed; pass them via params (ParamRef/params.xxx)."
+    )
+
+
+@dataclass(frozen=True, slots=True, init=False)
 class Expr:
-    """Expression tree node evaluated against a FormContext."""
+    """Expression tree node evaluated against a FormContext.
+
+    Compile flow (recommended):
+    - build an Expr via operators/refs
+    - compile_* builds an EvalPlan (postorder nodes + index)
+    - eval_with_plan(plan, ctx, params, u_elem) evaluates per element
+
+    Expr.eval is a debug/single-shot path that creates a plan on demand.
+    """
 
-    def __init__(self, op: str, *args):
+    op: OpName
+    args: tuple[Any, ...]
+
+    def __init__(self, op: OpName, *args):
+        if op not in _OP_NAMES:
+            raise ValueError(f"Unknown Expr op: {op!r}")
         object.__setattr__(self, "op", op)
         object.__setattr__(self, "args", args)
 
     def eval(self, ctx, params=None, u_elem=None):
+        """Evaluate the expression against a context (debug/single-shot path)."""
         return _eval_expr(self, ctx, params, u_elem=u_elem)
 
+    def children(self) -> tuple[Any, ...]:
+        """Return direct child nodes (Expr/FieldRef/ParamRef) for traversal."""
+        return tuple(arg for arg in self.args if isinstance(arg, (Expr, FieldRef, ParamRef)))
+
+    def walk(self) -> Iterator[Any]:
+        """Depth-first walk over nodes, including leaf FieldRef/ParamRef."""
+        yield self
+        for child in self.children():
+            if isinstance(child, Expr):
+                yield from child.walk()
+            else:
+                yield child
+
+    def postorder(self) -> Iterator[Any]:
+        """Postorder walk over nodes, including leaf FieldRef/ParamRef."""
+        for child in self.children():
+            if isinstance(child, Expr):
+                yield from child.postorder()
+            else:
+                yield child
+        yield self
+
+    def postorder_expr(self) -> Iterator["Expr"]:
+        """Postorder walk over Expr nodes only (for eval planning)."""
+        for arg in self.args:
+            if isinstance(arg, Expr):
+                yield from arg.postorder_expr()
+        yield self
+
     def _binop(self, other, op):
         return Expr(op, self, _as_expr(other))
 
     def __add__(self, other):
+        """Add expressions: `a + b`."""
         return self._binop(other, "add")
 
     def __radd__(self, other):
-        return _as_expr(other)._binop(self, "add")
+        """Right-add expressions: `1 + expr`."""
+        return Expr("add", _as_expr(other), self)
 
     def __sub__(self, other):
+        """Subtract expressions: `a - b`."""
         return self._binop(other, "sub")
 
     def __rsub__(self, other):
-        return _as_expr(other)._binop(self, "sub")
+        """Right-subtract expressions: `1 - expr`."""
+        return Expr("sub", _as_expr(other), self)
 
     def __mul__(self, other):
+        """Multiply expressions: `a * b`."""
         return self._binop(other, "mul")
 
     def __rmul__(self, other):
-        return _as_expr(other)._binop(self, "mul")
+        """Right-multiply expressions: `2 * expr`."""
+        return Expr("mul", _as_expr(other), self)
 
     def __matmul__(self, other):
+        """Matrix product: `a @ b` (FEM-specific contraction semantics)."""
         return self._binop(other, "matmul")
 
     def __rmatmul__(self, other):
-        return _as_expr(other)._binop(self, "matmul")
+        """Right-matmul: `A @ expr`."""
+        return Expr("matmul", _as_expr(other), self)
 
     def __or__(self, other):
-        return self._binop(other, "inner")
+        """Tensor inner product: `a | b` (use .val/.grad for FieldRef)."""
+        if isinstance(other, FieldRef):
+            raise TypeError("FieldRef | FieldRef is not supported; use outer(test, trial).")
+        return Expr("inner", self, _as_expr(other))
 
     def __ror__(self, other):
-        return _as_expr(other)._binop(self, "inner")
+        """Tensor inner product: `a | b` (use .val/.grad for FieldRef)."""
+        if isinstance(other, FieldRef):
+            raise TypeError("FieldRef | FieldRef is not supported; use outer(test, trial).")
+        return Expr("inner", _as_expr(other), self)
 
     def __pow__(self, power, modulo=None):
+        """Power: `a ** p` (no modulo support)."""
         if modulo is not None:
             raise ValueError("modulo is not supported for Expr exponentiation.")
         return Expr("pow", self, _as_expr(power))
 
     def __neg__(self):
+        """Unary negation: `-expr`."""
         return Expr("neg", self)
 
     @property
     def T(self):
+        """Transpose view: `expr.T`."""
         return Expr("transpose", self)
 
+    def __repr__(self) -> str:
+        return self.pretty()
 
-@dataclass(frozen=True)
-class FieldRef(Expr):
+    def pretty(self) -> str:
+        return _pretty_expr(self)
+
+
+@dataclass(frozen=True, slots=True)
+class FieldRef:
     """Symbolic reference to trial/test/unknown field, optionally by name."""
 
     role: str
     name: str | None = None
 
-    def __init__(self, role: str, name: str | None = None):
-        object.__setattr__(self, "role", role)
-        object.__setattr__(self, "name", name)
-        super().__init__("field", role, name)
-
     @property
     def val(self):
         return Expr("value", self)
@@ -91,12 +321,18 @@ class FieldRef(Expr):
 
     def __mul__(self, other):
         if isinstance(other, FieldRef):
-            return Expr("outer", self, other)
+            raise TypeError(
+                "FieldRef * FieldRef is ambiguous; use outer(v, u) (test, trial), "
+                "action(v, s), or dot(v, q)."
+            )
         return Expr("mul", Expr("value", self), _as_expr(other))
 
     def __rmul__(self, other):
         if isinstance(other, FieldRef):
-            return Expr("outer", other, self)
+            raise TypeError(
+                "FieldRef * FieldRef is ambiguous; use outer(v, u) (test, trial), "
+                "action(v, s), or dot(v, q)."
+            )
         return Expr("mul", _as_expr(other), Expr("value", self))
 
     def __add__(self, other):
@@ -113,22 +349,23 @@ class FieldRef(Expr):
 
     def __or__(self, other):
         if isinstance(other, FieldRef):
-            return Expr("inner", self, other)
-        return Expr("sdot", self, _as_expr(other))
+            raise TypeError(
+                "FieldRef | FieldRef is not supported; use outer(test, trial) for basis kernels."
+            )
+        return Expr("dot", self, _as_expr(other))
 
     def __ror__(self, other):
         if isinstance(other, FieldRef):
-            return Expr("inner", other, self)
-        return Expr("sdot", _as_expr(other), self)
+            raise TypeError(
+                "FieldRef | FieldRef is not supported; use outer(test, trial) for basis kernels."
+            )
+        return Expr("dot", _as_expr(other), self)
 
 
-@dataclass(frozen=True)
-class ParamRef(Expr):
+@dataclass(frozen=True, slots=True)
+class ParamRef:
     """Symbolic reference to params passed into the kernel."""
 
-    def __init__(self):
-        super().__init__("param")
-
     def __getattr__(self, name: str):
         return Expr("getattr", self, name)
 
@@ -179,13 +416,11 @@ def param_ref() -> ParamRef:
     return ParamRef()
 
 
-def _as_expr(obj) -> Expr:
-    if isinstance(obj, Expr):
-        return obj
-    return Expr("lit", obj)
-
-
-def _eval_field(obj: Any, ctx, params):
+def _eval_field(
+    obj: Any,
+    ctx: VolumeContext | SurfaceContext,
+    params: ParamsLike,
+) -> FormFieldLike:
     if isinstance(obj, FieldRef):
         if obj.name is not None:
             mixed_fields = getattr(ctx, "fields", None)
@@ -226,25 +461,23 @@ def _eval_field(obj: Any, ctx, params):
         if obj.role == "unknown":
             return getattr(ctx, "unknown", ctx.trial)
         raise ValueError(f"Unknown field role: {obj.role}")
-    if isinstance(obj, Expr):
-        val = obj.eval(ctx, params)
-        if hasattr(val, "N"):
-            return val
     raise TypeError("Expected a field reference for this operator.")
 
 
-def _eval_value(obj: Any, ctx, params, u_elem=None):
-    if isinstance(obj, FieldRef):
-        field = _eval_field(obj, ctx, params)
-        if obj.role == "unknown":
-            return _eval_unknown_value(obj, field, u_elem)
-        return field.N
-    if isinstance(obj, Expr):
-        return obj.eval(ctx, params, u_elem=u_elem)
-    return obj
+# def _eval_value(obj: Any, ctx, params, u_elem=None):
+#     if isinstance(obj, FieldRef):
+#         field = _eval_field(obj, ctx, params)
+#         if obj.role == "unknown":
+#             return _eval_unknown_value(obj, field, u_elem)
+#         return field.N
+#     if isinstance(obj, ParamRef):
+#         return params
+#     if isinstance(obj, Expr):
+#         return obj.eval(ctx, params, u_elem=u_elem)
+#     return obj
 
 
-def _extract_unknown_elem(field_ref: FieldRef, u_elem):
+def _extract_unknown_elem(field_ref: FieldRef, u_elem: UElement):
     if u_elem is None:
         raise ValueError("u_elem is required to evaluate unknown field value.")
     if isinstance(u_elem, dict):
@@ -255,7 +488,17 @@ def _extract_unknown_elem(field_ref: FieldRef, u_elem):
     return u_elem
 
 
-def _eval_unknown_value(field_ref: FieldRef, field, u_elem):
+def _basis_outer(test: FieldRef, trial: FieldRef, ctx, params):
+    v_field = _eval_field(test, ctx, params)
+    u_field = _eval_field(trial, ctx, params)
+    if getattr(v_field, "value_dim", 1) != 1 or getattr(u_field, "value_dim", 1) != 1:
+        raise ValueError(
+            "inner/outer is only defined for scalar fields; use dot/action/einsum for vector/tensor cases."
+        )
+    return jnp.einsum("qi,qj->qij", v_field.N, u_field.N)
+
+
+def _eval_unknown_value(field_ref: FieldRef, field: FormFieldLike, u_elem: UElement):
     u_local = _extract_unknown_elem(field_ref, u_elem)
     value_dim = int(getattr(field, "value_dim", 1))
     if value_dim == 1:
@@ -264,7 +507,7 @@ def _eval_unknown_value(field_ref: FieldRef, field, u_elem):
     return jnp.einsum("qa,ai->qi", field.N, u_nodes)
 
 
-def _eval_unknown_grad(field_ref: FieldRef, field, u_elem):
+def _eval_unknown_grad(field_ref: FieldRef, field: FormFieldLike, u_elem: UElement):
     u_local = _extract_unknown_elem(field_ref, u_elem)
     if u_local is None:
         raise ValueError("u_elem is required to evaluate unknown field gradient.")
@@ -285,6 +528,15 @@ def sym_grad(field) -> Expr:
     return Expr("sym_grad", _as_expr(field))
 
 
+def outer(a, b) -> Expr:
+    """Outer product of scalar fields: `outer(v, u)` (test, trial)."""
+    if not isinstance(a, FieldRef) or not isinstance(b, FieldRef):
+        raise TypeError("outer expects FieldRef operands.")
+    if a.role != "test" or b.role != "trial":
+        raise TypeError("outer expects outer(test, trial).")
+    return Expr("outer", a, b)
+
+
 def dot(a, b) -> Expr:
     """Dot product or vector load helper."""
     return Expr("dot", _as_expr(a), _as_expr(b))
@@ -303,7 +555,7 @@ def ddot(a, b, c=None) -> Expr:
 
 
 def inner(a, b) -> Expr:
-    """Inner product over the last axis."""
+    """Inner product over the last axis (tensor-level)."""
     return Expr("inner", _as_expr(a), _as_expr(b))
 
 
@@ -362,6 +614,16 @@ def transpose_last2(a) -> Expr:
     return Expr("transpose_last2", _as_expr(a))
 
 
+def matmul(a, b) -> Expr:
+    """FEM-specific batched contraction (same semantics as `@`)."""
+    return Expr("matmul", _as_expr(a), _as_expr(b))
+
+
+def matmul_std(a, b) -> Expr:
+    """Standard matrix product (`jnp.matmul` semantics)."""
+    return Expr("matmul_std", _as_expr(a), _as_expr(b))
+
+
 def einsum(subscripts: str, *args) -> Expr:
     """Einsum wrapper that supports Expr inputs."""
     return Expr("einsum", subscripts, *[_as_expr(arg) for arg in args])
@@ -369,9 +631,22 @@ def einsum(subscripts: str, *args) -> Expr:
 
 def _call_user(fn, *args, params):
     try:
+        sig = inspect.signature(fn)
+    except (TypeError, ValueError):
+        return fn(*args, params)
+
+    params_list = list(sig.parameters.values())
+    if any(p.kind == inspect.Parameter.VAR_POSITIONAL for p in params_list):
+        return fn(*args, params)
+    positional = [
+        p
+        for p in params_list
+        if p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+    ]
+    max_positional = len(positional)
+    if len(args) + 1 <= max_positional:
         return fn(*args, params)
-    except TypeError:
-        return fn(*args)
+    return fn(*args)
 
 
 def compile_bilinear(fn):
@@ -382,19 +657,26 @@ def compile_bilinear(fn):
         u = trial_ref()
         v = test_ref()
         p = param_ref()
-        try:
-            expr = fn(u, v, p)
-        except TypeError:
-            expr = fn(u, v)
+        expr = _call_user(fn, u, v, params=p)
+    expr = _as_expr(expr)
+    if not isinstance(expr, Expr):
+        raise TypeError("Bilinear form must return an Expr.")
 
-    includes_measure = _expr_contains(expr, "volume_measure")
-    if not includes_measure:
+    volume_count = _count_op(expr, "volume_measure")
+    surface_count = _count_op(expr, "surface_measure")
+    if volume_count == 0:
         raise ValueError("Volume bilinear form must include dOmega().")
+    if volume_count > 1:
+        raise ValueError("Volume bilinear form must include dOmega() exactly once.")
+    if surface_count > 0:
+        raise ValueError("Volume bilinear form must not include ds().")
+
+    plan = make_eval_plan(expr)
 
     def _form(ctx, params):
-        return _as_expr(expr).eval(ctx, params)
+        return eval_with_plan(plan, ctx, params)
 
-    _form._includes_measure = includes_measure
+    _form._includes_measure = True
     return _form
 
 
@@ -405,19 +687,26 @@ def compile_linear(fn):
     else:
         v = test_ref()
         p = param_ref()
-        try:
-            expr = fn(v, p)
-        except TypeError:
-            expr = fn(v)
+        expr = _call_user(fn, v, params=p)
+    expr = _as_expr(expr)
+    if not isinstance(expr, Expr):
+        raise TypeError("Linear form must return an Expr.")
 
-    includes_measure = _expr_contains(expr, "volume_measure")
-    if not includes_measure:
+    volume_count = _count_op(expr, "volume_measure")
+    surface_count = _count_op(expr, "surface_measure")
+    if volume_count == 0:
         raise ValueError("Volume linear form must include dOmega().")
+    if volume_count > 1:
+        raise ValueError("Volume linear form must include dOmega() exactly once.")
+    if surface_count > 0:
+        raise ValueError("Volume linear form must not include ds().")
+
+    plan = make_eval_plan(expr)
 
     def _form(ctx, params):
-        return _as_expr(expr).eval(ctx, params)
+        return eval_with_plan(plan, ctx, params)
 
-    _form._includes_measure = includes_measure
+    _form._includes_measure = True
     return _form
 
 
@@ -429,6 +718,340 @@ def _expr_contains(expr: Expr, op: str) -> bool:
     return any(_expr_contains(arg, op) for arg in expr.args if isinstance(arg, Expr))
 
 
+def _count_op(expr: Expr, op: str) -> int:
+    if not isinstance(expr, Expr):
+        return 0
+    count = 1 if expr.op == op else 0
+    for arg in expr.args:
+        if isinstance(arg, Expr):
+            count += _count_op(arg, op)
+    return count
+
+
+@dataclass(frozen=True, slots=True)
+class EvalPlan:
+    expr: Expr
+    nodes: tuple[Expr, ...]
+    index: dict[int, int]
+
+
+def _validate_eval_plan(nodes: tuple[Expr, ...]) -> None:
+    fieldref_arg_ops = {
+        "value",
+        "grad",
+        "sym_grad",
+        "dot",
+        "sdot",
+        "action",
+        "gaction",
+        "outer",
+    }
+    for node in nodes:
+        op = node.op
+        args = node.args
+        if op not in fieldref_arg_ops:
+            if any(isinstance(arg, FieldRef) for arg in args):
+                raise TypeError(f"{op} cannot take FieldRef directly; wrap with .val/.grad/.sym_grad.")
+        if op in {"value", "grad", "sym_grad"}:
+            if len(args) != 1 or not isinstance(args[0], FieldRef):
+                raise TypeError(f"{op} expects FieldRef.")
+        elif op in {"dot", "sdot"}:
+            if len(args) != 2:
+                raise TypeError(f"{op} expects two arguments.")
+            if any(isinstance(arg, FieldRef) for arg in args):
+                if not isinstance(args[0], FieldRef):
+                    raise TypeError(f"{op} expects FieldRef as the first argument.")
+                if isinstance(args[1], FieldRef):
+                    raise TypeError(f"{op} expects an expression for the second argument; use .val/.grad.")
+        elif op in {"action", "gaction"}:
+            if len(args) != 2 or not isinstance(args[0], FieldRef):
+                raise TypeError(f"{op} expects FieldRef as the first argument.")
+            if op == "action" and isinstance(args[1], FieldRef):
+                raise ValueError("action expects a scalar expression; use u.val for unknowns.")
+            if op == "gaction" and isinstance(args[1], FieldRef):
+                raise TypeError("gaction expects an expression for the second argument; use .grad.")
+        elif op == "outer":
+            if len(args) != 2 or not all(isinstance(arg, FieldRef) for arg in args):
+                raise TypeError("outer expects two FieldRef operands.")
+            if args[0].role != "test" or args[1].role != "trial":
+                raise TypeError("outer expects outer(test, trial).")
+
+
+def make_eval_plan(expr: Expr) -> EvalPlan:
+    nodes = tuple(expr.postorder_expr())
+    _validate_eval_plan(nodes)
+    index: dict[int, int] = {}
+    for i, node in enumerate(nodes):
+        index.setdefault(id(node), i)
+    return EvalPlan(expr=expr, nodes=nodes, index=index)
+
+
+def eval_with_plan(
+    plan: EvalPlan,
+    ctx: VolumeContext | SurfaceContext,
+    params: ParamsLike,
+    u_elem: UElement | None = None,
+):
+    nodes = plan.nodes
+    index = plan.index
+    vals: list[Any] = [None] * len(nodes)
+
+    def get(obj):
+        if isinstance(obj, Expr):
+            return vals[index[id(obj)]]
+        if isinstance(obj, FieldRef):
+            raise TypeError(
+                "FieldRef must be wrapped with .val/.grad/.sym_grad or used as the first arg of dot/action."
+            )
+        if isinstance(obj, ParamRef):
+            return params
+        return obj
+
+    for i, node in enumerate(nodes):
+        op = node.op
+        args = node.args
+
+        if op == "lit":
+            vals[i] = args[0]
+            continue
+        if op == "getattr":
+            base = get(args[0])
+            name = args[1]
+            if isinstance(base, dict):
+                vals[i] = base[name]
+            else:
+                vals[i] = getattr(base, name)
+            continue
+        if op == "value":
+            ref = args[0]
+            assert isinstance(ref, FieldRef)
+            field = _eval_field(ref, ctx, params)
+            if ref.role == "unknown":
+                vals[i] = _eval_unknown_value(ref, field, u_elem)
+            else:
+                vals[i] = field.N
+            continue
+        if op == "grad":
+            ref = args[0]
+            assert isinstance(ref, FieldRef)
+            field = _eval_field(ref, ctx, params)
+            if ref.role == "unknown":
+                vals[i] = _eval_unknown_grad(ref, field, u_elem)
+            else:
+                vals[i] = field.gradN
+            continue
+        if op == "pow":
+            base = get(args[0])
+            exp = get(args[1])
+            vals[i] = base**exp
+            continue
+        if op == "eye":
+            vals[i] = jnp.eye(int(args[0]))
+            continue
+        if op == "det":
+            vals[i] = jnp.linalg.det(get(args[0]))
+            continue
+        if op == "inv":
+            vals[i] = jnp.linalg.inv(get(args[0]))
+            continue
+        if op == "transpose":
+            vals[i] = jnp.swapaxes(get(args[0]), -1, -2)
+            continue
+        if op == "log":
+            vals[i] = jnp.log(get(args[0]))
+            continue
+        if op == "surface_normal":
+            normal = getattr(ctx, "normal", None)
+            if normal is None:
+                raise ValueError("surface normal is not available in context")
+            vals[i] = normal
+            continue
+        if op == "surface_measure":
+            if not hasattr(ctx, "w") or not hasattr(ctx, "detJ"):
+                raise TypeError("surface measure requires SurfaceContext.")
+            vals[i] = ctx.w * ctx.detJ
+            continue
+        if op == "volume_measure":
+            if not hasattr(ctx, "w") or not hasattr(ctx, "test"):
+                raise TypeError("volume measure requires VolumeContext.")
+            vals[i] = ctx.w * ctx.test.detJ
+            continue
+        if op == "sym_grad":
+            ref = args[0]
+            assert isinstance(ref, FieldRef)
+            field = _eval_field(ref, ctx, params)
+            if ref.role == "unknown":
+                if u_elem is None:
+                    raise ValueError("u_elem is required to evaluate unknown sym_grad.")
+                u_local = _extract_unknown_elem(ref, u_elem)
+                vals[i] = _ops.sym_grad_u(field, u_local)
+            else:
+                vals[i] = _ops.sym_grad(field)
+            continue
+        if op == "outer":
+            a, b = args
+            if not isinstance(a, FieldRef) or not isinstance(b, FieldRef):
+                raise TypeError("outer expects FieldRef operands.")
+            test, trial = a, b
+            vals[i] = _basis_outer(test, trial, ctx, params)
+            continue
+        if op == "add":
+            vals[i] = get(args[0]) + get(args[1])
+            continue
+        if op == "sub":
+            vals[i] = get(args[0]) - get(args[1])
+            continue
+        if op == "mul":
+            a = get(args[0])
+            b = get(args[1])
+            if hasattr(a, "ndim") and hasattr(b, "ndim"):
+                if a.ndim == 1 and b.ndim == 2 and a.shape[0] == b.shape[0]:
+                    a = a[:, None]
+                elif b.ndim == 1 and a.ndim == 2 and b.shape[0] == a.shape[0]:
+                    b = b[:, None]
+                elif a.ndim >= 2 and b.ndim == 1 and a.shape[0] == b.shape[0]:
+                    b = b.reshape((b.shape[0],) + (1,) * (a.ndim - 1))
+                elif b.ndim >= 2 and a.ndim == 1 and b.shape[0] == a.shape[0]:
+                    a = a.reshape((a.shape[0],) + (1,) * (b.ndim - 1))
+            vals[i] = a * b
+            continue
+        if op == "matmul":
+            a = get(args[0])
+            b = get(args[1])
+            if (
+                hasattr(a, "ndim")
+                and hasattr(b, "ndim")
+                and a.ndim == 3
+                and b.ndim == 3
+                and a.shape[0] == b.shape[0]
+                and a.shape[-1] == b.shape[-1]
+            ):
+                vals[i] = jnp.einsum("qia,qja->qij", a, b)
+            else:
+                raise TypeError(
+                    "Expr '@' (matmul) is FEM-specific; use matmul_std(a, b) for standard matmul."
+                )
+            continue
+        if op == "matmul_std":
+            a = get(args[0])
+            b = get(args[1])
+            vals[i] = jnp.matmul(a, b)
+            continue
+        if op == "neg":
+            vals[i] = -get(args[0])
+            continue
+        if op == "dot":
+            ref = args[0]
+            if isinstance(ref, FieldRef):
+                vals[i] = _ops.dot(_eval_field(ref, ctx, params), get(args[1]))
+            else:
+                a = get(args[0])
+                b = get(args[1])
+                if (
+                    hasattr(a, "ndim")
+                    and hasattr(b, "ndim")
+                    and a.ndim == 3
+                    and b.ndim == 3
+                    and a.shape[-1] == b.shape[-1]
+                ):
+                    vals[i] = jnp.einsum("qia,qja->qij", a, b)
+                else:
+                    vals[i] = jnp.matmul(a, b)
+            continue
+        if op == "sdot":
+            ref = args[0]
+            if isinstance(ref, FieldRef):
+                vals[i] = _ops.dot(_eval_field(ref, ctx, params), get(args[1]))
+            else:
+                a = get(args[0])
+                b = get(args[1])
+                if (
+                    hasattr(a, "ndim")
+                    and hasattr(b, "ndim")
+                    and a.ndim == 3
+                    and b.ndim == 3
+                    and a.shape[-1] == b.shape[-1]
+                ):
+                    vals[i] = jnp.einsum("qia,qja->qij", a, b)
+                else:
+                    vals[i] = jnp.matmul(a, b)
+            continue
+        if op == "ddot":
+            if len(args) == 2:
+                a = get(args[0])
+                b = get(args[1])
+                if (
+                    hasattr(a, "ndim")
+                    and hasattr(b, "ndim")
+                    and a.ndim == 3
+                    and b.ndim == 3
+                    and a.shape[0] == b.shape[0]
+                    and a.shape[1] == b.shape[1]
+                ):
+                    vals[i] = jnp.einsum("qik,qim->qkm", a, b)
+                else:
+                    vals[i] = _ops.ddot(a, b)
+            else:
+                vals[i] = _ops.ddot(get(args[0]), get(args[1]), get(args[2]))
+            continue
+        if op == "inner":
+            a = get(args[0])
+            b = get(args[1])
+            vals[i] = jnp.einsum("...i,...i->...", a, b)
+            continue
+        if op == "action":
+            ref = args[0]
+            assert isinstance(ref, FieldRef)
+            if isinstance(args[1], FieldRef):
+                raise ValueError("action expects a scalar expression; use u.val for unknowns.")
+            v_field = _eval_field(ref, ctx, params)
+            s = get(args[1])
+            value_dim = int(getattr(v_field, "value_dim", 1))
+            # action maps a test field with a scalar/vector expression into nodal space.
+            if value_dim == 1:
+                if v_field.N.ndim != 2:
+                    raise ValueError("action expects scalar test field with N shape (q, ndofs).")
+                if hasattr(s, "ndim") and s.ndim not in (0, 1):
+                    raise ValueError("action expects scalar s with shape (q,) or scalar.")
+                vals[i] = v_field.N * s
+            else:
+                if hasattr(s, "ndim") and s.ndim not in (1, 2):
+                    raise ValueError("action expects vector s with shape (q, dim) or (dim,).")
+                vals[i] = _ops.dot(v_field, s)
+            continue
+        if op == "gaction":
+            ref = args[0]
+            assert isinstance(ref, FieldRef)
+            v_field = _eval_field(ref, ctx, params)
+            q = get(args[1])
+            # gaction maps a flux-like expression to nodal space via test gradients.
+            if v_field.gradN.ndim != 3:
+                raise ValueError("gaction expects test gradient with shape (q, ndofs, dim).")
+            if not hasattr(q, "ndim"):
+                raise ValueError("gaction expects q with shape (q, dim) or (q, dim, dim).")
+            if q.ndim == 2:
+                vals[i] = jnp.einsum("qaj,qj->qa", v_field.gradN, q)
+            elif q.ndim == 3:
+                if int(getattr(v_field, "value_dim", 1)) == 1:
+                    raise ValueError("gaction tensor flux requires vector test field.")
+                vals[i] = jnp.einsum("qij,qaj->qai", q, v_field.gradN).reshape(q.shape[0], -1)
+            else:
+                raise ValueError("gaction expects q with shape (q, dim) or (q, dim, dim).")
+            continue
+        if op == "transpose_last2":
+            vals[i] = _ops.transpose_last2(get(args[0]))
+            continue
+        if op == "einsum":
+            subscripts = args[0]
+            operands = [get(arg) for arg in args[1:]]
+            vals[i] = jnp.einsum(subscripts, *operands)
+            continue
+
+        raise ValueError(f"Unknown Expr op: {op}")
+
+    return vals[index[id(plan.expr)]]
+
+
 def compile_surface_linear(fn):
     """get_compiled a surface linear form into a kernel (ctx, params) -> ndarray."""
     if isinstance(fn, Expr):
@@ -436,26 +1059,27 @@ def compile_surface_linear(fn):
     else:
         v = test_ref()
         p = param_ref()
-        expr = None
-        try:
-            expr = fn(v, p)
-        except TypeError:
-            try:
-                expr = fn(v)
-            except TypeError:
-                expr = None
+        expr = _call_user(fn, v, params=p)
 
+    expr = _as_expr(expr)
     if not isinstance(expr, Expr):
         raise ValueError("Surface linear form must return an Expr; use ds() in the expression.")
 
-    includes_measure = _expr_contains(expr, "surface_measure")
-    if not includes_measure:
+    surface_count = _count_op(expr, "surface_measure")
+    volume_count = _count_op(expr, "volume_measure")
+    if surface_count == 0:
         raise ValueError("Surface linear form must include ds().")
+    if surface_count > 1:
+        raise ValueError("Surface linear form must include ds() exactly once.")
+    if volume_count > 0:
+        raise ValueError("Surface linear form must not include dOmega().")
+
+    plan = make_eval_plan(expr)
 
     def _form(ctx, params):
-        return _as_expr(expr).eval(ctx, params)
+        return eval_with_plan(plan, ctx, params)
 
-    _form._includes_measure = includes_measure  # type: ignore[attr-defined]
+    _form._includes_measure = True  # type: ignore[attr-defined]
     return _form
 
 
@@ -519,25 +1143,33 @@ def compile_residual(fn):
         v = test_ref()
         u = unknown_ref()
         p = param_ref()
-        try:
-            expr = fn(v, u, p)
-        except TypeError:
-            expr = fn(v, u)
+        expr = _call_user(fn, v, u, params=p)
+    expr = _as_expr(expr)
+    if not isinstance(expr, Expr):
+        raise TypeError("Residual form must return an Expr.")
 
-    includes_measure = _expr_contains(expr, "volume_measure")
-    if not includes_measure:
+    volume_count = _count_op(expr, "volume_measure")
+    surface_count = _count_op(expr, "surface_measure")
+    if volume_count == 0:
         raise ValueError("Volume residual form must include dOmega().")
+    if volume_count > 1:
+        raise ValueError("Volume residual form must include dOmega() exactly once.")
+    if surface_count > 0:
+        raise ValueError("Volume residual form must not include ds().")
+
+    plan = make_eval_plan(expr)
 
     def _form(ctx, u_elem, params):
-        return _as_expr(expr).eval(ctx, params, u_elem=u_elem)
+        return eval_with_plan(plan, ctx, params, u_elem=u_elem)
 
-    _form._includes_measure = includes_measure
+    _form._includes_measure = True
     return _form
 
 
 def compile_mixed_residual(residuals: dict[str, Callable]):
     """get_compiled mixed residuals keyed by field name."""
     compiled = {}
+    plans = {}
     includes_measure = {}
     for name, fn in residuals.items():
         if isinstance(fn, Expr):
@@ -546,17 +1178,24 @@ def compile_mixed_residual(residuals: dict[str, Callable]):
             v = test_ref(name)
             u = unknown_ref(name)
             p = param_ref()
-            try:
-                expr = fn(v, u, p)
-            except TypeError:
-                expr = fn(v, u)
-        compiled[name] = _as_expr(expr)
-        includes_measure[name] = _expr_contains(compiled[name], "volume_measure")
-        if not includes_measure[name]:
+            expr = _call_user(fn, v, u, params=p)
+        expr = _as_expr(expr)
+        if not isinstance(expr, Expr):
+            raise TypeError(f"Mixed residual '{name}' must return an Expr.")
+        compiled[name] = expr
+        plans[name] = make_eval_plan(expr)
+        volume_count = _count_op(compiled[name], "volume_measure")
+        surface_count = _count_op(compiled[name], "surface_measure")
+        includes_measure[name] = volume_count == 1
+        if volume_count == 0:
             raise ValueError(f"Mixed residual '{name}' must include dOmega().")
+        if volume_count > 1:
+            raise ValueError(f"Mixed residual '{name}' must include dOmega() exactly once.")
+        if surface_count > 0:
+            raise ValueError(f"Mixed residual '{name}' must not include ds().")
 
     def _form(ctx, u_elem, params):
-        return {name: expr.eval(ctx, params, u_elem=u_elem) for name, expr in compiled.items()}
+        return {name: eval_with_plan(plan, ctx, params, u_elem=u_elem) for name, plan in plans.items()}
 
     _form._includes_measure = includes_measure
     return _form
@@ -574,216 +1213,14 @@ class MixedWeakForm:
         return compile_mixed_residual(self.residuals)
 
 
-def _eval_expr(expr: Expr, ctx, params, u_elem=None):
-    op = expr.op
-    args = expr.args
-
-    if op == "lit":
-        return args[0]
-    if op == "param":
-        return params
-    if op == "getattr":
-        base = _eval_value(args[0], ctx, params, u_elem=u_elem)
-        name = args[1]
-        if isinstance(base, dict):
-            return base[name]
-        return getattr(base, name)
-    if op == "field":
-        role, name = args
-        if name is not None:
-            if role == "trial" and getattr(ctx, "trial_fields", None) is not None:
-                if name in ctx.trial_fields:
-                    return ctx.trial_fields[name]
-            if role == "test" and getattr(ctx, "test_fields", None) is not None:
-                if name in ctx.test_fields:
-                    return ctx.test_fields[name]
-            if role == "unknown" and getattr(ctx, "unknown_fields", None) is not None:
-                if name in ctx.unknown_fields:
-                    return ctx.unknown_fields[name]
-            fields = getattr(ctx, "fields", None)
-            if fields is not None and name in fields:
-                group = fields[name]
-                if isinstance(group, dict):
-                    if role in group:
-                        return group[role]
-                    if "field" in group:
-                        return group["field"]
-                return group
-        if role == "trial":
-            return ctx.trial
-        if role == "test":
-            return ctx.test
-        if role == "unknown":
-            return getattr(ctx, "unknown", ctx.trial)
-        raise ValueError(f"Unknown field role: {role}")
-    if op == "value":
-        field = _eval_field(args[0], ctx, params)
-        if isinstance(args[0], FieldRef) and args[0].role == "unknown":
-            return _eval_unknown_value(args[0], field, u_elem)
-        return field.N
-    if op == "grad":
-        field = _eval_field(args[0], ctx, params)
-        if isinstance(args[0], FieldRef) and args[0].role == "unknown":
-            return _eval_unknown_grad(args[0], field, u_elem)
-        return field.gradN
-    if op == "pow":
-        base = _eval_value(args[0], ctx, params, u_elem=u_elem)
-        exp = _eval_value(args[1], ctx, params, u_elem=u_elem)
-        return base**exp
-    if op == "eye":
-        return jnp.eye(int(args[0]))
-    if op == "det":
-        return jnp.linalg.det(_eval_value(args[0], ctx, params, u_elem=u_elem))
-    if op == "inv":
-        return jnp.linalg.inv(_eval_value(args[0], ctx, params, u_elem=u_elem))
-    if op == "transpose":
-        return jnp.swapaxes(_eval_value(args[0], ctx, params, u_elem=u_elem), -1, -2)
-    if op == "log":
-        return jnp.log(_eval_value(args[0], ctx, params, u_elem=u_elem))
-    if op == "surface_normal":
-        normal = getattr(ctx, "normal", None)
-        if normal is None:
-            raise ValueError("surface normal is not available in context")
-        return normal
-    if op == "surface_measure":
-        if not hasattr(ctx, "w") or not hasattr(ctx, "detJ"):
-            raise ValueError("surface measure requires surface context with w and detJ.")
-        return ctx.w * ctx.detJ
-    if op == "volume_measure":
-        if not hasattr(ctx, "w") or not hasattr(ctx, "test"):
-            raise ValueError("volume measure requires FormContext with w and test.detJ.")
-        return ctx.w * ctx.test.detJ
-    if op == "sym_grad":
-        field = _eval_field(args[0], ctx, params)
-        if isinstance(args[0], FieldRef) and args[0].role == "unknown":
-            if u_elem is None:
-                raise ValueError("u_elem is required to evaluate unknown sym_grad.")
-            u_local = _extract_unknown_elem(args[0], u_elem)
-            return _ops.sym_grad_u(field, u_local)
-        return _ops.sym_grad(field)
-    if op == "outer":
-        a, b = args
-        if not isinstance(a, FieldRef) or not isinstance(b, FieldRef):
-            raise TypeError("outer expects FieldRef operands.")
-        if a.role == b.role:
-            raise ValueError("outer requires one trial and one test field.")
-        test = a if a.role == "test" else b
-        trial = b if a.role == "test" else a
-        v_field = _eval_field(test, ctx, params)
-        u_field = _eval_field(trial, ctx, params)
-        if getattr(v_field, "value_dim", 1) != 1 or getattr(u_field, "value_dim", 1) != 1:
-            raise ValueError("u*v is only defined for scalar fields; use dot/inner for vectors.")
-        vN = v_field.N
-        uN = u_field.N
-        return jnp.einsum("qi,qj->qij", vN, uN)
-    if op == "add":
-        return _eval_value(args[0], ctx, params, u_elem=u_elem) + _eval_value(args[1], ctx, params, u_elem=u_elem)
-    if op == "sub":
-        return _eval_value(args[0], ctx, params, u_elem=u_elem) - _eval_value(args[1], ctx, params, u_elem=u_elem)
-    if op == "mul":
-        a = _eval_value(args[0], ctx, params, u_elem=u_elem)
-        b = _eval_value(args[1], ctx, params, u_elem=u_elem)
-        if hasattr(a, "ndim") and hasattr(b, "ndim"):
-            if a.ndim == 1 and b.ndim == 2 and a.shape[0] == b.shape[0]:
-                a = a[:, None]
-            elif b.ndim == 1 and a.ndim == 2 and b.shape[0] == a.shape[0]:
-                b = b[:, None]
-            elif a.ndim >= 2 and b.ndim == 1 and a.shape[0] == b.shape[0]:
-                b = b.reshape((b.shape[0],) + (1,) * (a.ndim - 1))
-            elif b.ndim >= 2 and a.ndim == 1 and b.shape[0] == a.shape[0]:
-                a = a.reshape((a.shape[0],) + (1,) * (b.ndim - 1))
-        return a * b
-    if op == "matmul":
-        a = _eval_value(args[0], ctx, params, u_elem=u_elem)
-        b = _eval_value(args[1], ctx, params, u_elem=u_elem)
-        if (
-            hasattr(a, "ndim")
-            and hasattr(b, "ndim")
-            and a.ndim == 3
-            and b.ndim == 3
-            and a.shape[0] == b.shape[0]
-            and a.shape[-1] == b.shape[-1]
-        ):
-            return jnp.einsum("qia,qja->qij", a, b)
-        return a @ b
-    if op == "neg":
-        return -_eval_value(args[0], ctx, params, u_elem=u_elem)
-    if op == "dot":
-        if isinstance(args[0], FieldRef):
-            return _ops.dot(_eval_field(args[0], ctx, params), _eval_value(args[1], ctx, params, u_elem=u_elem))
-        a = _eval_value(args[0], ctx, params, u_elem=u_elem)
-        b = _eval_value(args[1], ctx, params, u_elem=u_elem)
-        if hasattr(a, "ndim") and hasattr(b, "ndim") and a.ndim == 3 and b.ndim == 3 and a.shape[-1] == b.shape[-1]:
-            return jnp.einsum("qia,qja->qij", a, b)
-        return jnp.matmul(a, b)
-    if op == "sdot":
-        if isinstance(args[0], FieldRef):
-            return _ops.dot(_eval_field(args[0], ctx, params), _eval_value(args[1], ctx, params, u_elem=u_elem))
-        a = _eval_value(args[0], ctx, params, u_elem=u_elem)
-        b = _eval_value(args[1], ctx, params, u_elem=u_elem)
-        if hasattr(a, "ndim") and hasattr(b, "ndim") and a.ndim == 3 and b.ndim == 3 and a.shape[-1] == b.shape[-1]:
-            return jnp.einsum("qia,qja->qij", a, b)
-        return jnp.matmul(a, b)
-    if op == "ddot":
-        if len(args) == 2:
-            a = _eval_value(args[0], ctx, params, u_elem=u_elem)
-            b = _eval_value(args[1], ctx, params, u_elem=u_elem)
-            if (
-                hasattr(a, "ndim")
-                and hasattr(b, "ndim")
-                and a.ndim == 3
-                and b.ndim == 3
-                and a.shape[0] == b.shape[0]
-                and a.shape[1] == b.shape[1]
-            ):
-                return jnp.einsum("qik,qim->qkm", a, b)
-            return _ops.ddot(a, b)
-        return _ops.ddot(
-            _eval_value(args[0], ctx, params, u_elem=u_elem),
-            _eval_value(args[1], ctx, params, u_elem=u_elem),
-            _eval_value(args[2], ctx, params, u_elem=u_elem),
-        )
-    if op == "inner":
-        a = _eval_value(args[0], ctx, params, u_elem=u_elem)
-        b = _eval_value(args[1], ctx, params, u_elem=u_elem)
-        return jnp.einsum("...i,...i->...", a, b)
-    if op == "action":
-        if isinstance(args[1], FieldRef):
-            raise ValueError("action expects a scalar expression; use u.val for unknowns.")
-        v_field = _eval_field(args[0], ctx, params)
-        s = _eval_value(args[1], ctx, params, u_elem=u_elem)
-        value_dim = int(getattr(v_field, "value_dim", 1))
-        if value_dim == 1:
-            if v_field.N.ndim != 2:
-                raise ValueError("action expects scalar test field with N shape (q, ndofs).")
-            if hasattr(s, "ndim") and s.ndim not in (0, 1):
-                raise ValueError("action expects scalar s with shape (q,) or scalar.")
-            return v_field.N * s
-        if hasattr(s, "ndim") and s.ndim not in (1, 2):
-            raise ValueError("action expects vector s with shape (q, dim) or (dim,).")
-        return _ops.dot(v_field, s)
-    if op == "gaction":
-        v_field = _eval_field(args[0], ctx, params)
-        q = _eval_value(args[1], ctx, params, u_elem=u_elem)
-        if v_field.gradN.ndim != 3:
-            raise ValueError("gaction expects test gradient with shape (q, ndofs, dim).")
-        if not hasattr(q, "ndim"):
-            raise ValueError("gaction expects q with shape (q, dim) or (q, dim, dim).")
-        if q.ndim == 2:
-            return jnp.einsum("qaj,qj->qa", v_field.gradN, q)
-        if q.ndim == 3:
-            if int(getattr(v_field, "value_dim", 1)) == 1:
-                raise ValueError("gaction tensor flux requires vector test field.")
-            return jnp.einsum("qij,qaj->qai", q, v_field.gradN).reshape(q.shape[0], -1)
-        raise ValueError("gaction expects q with shape (q, dim) or (q, dim, dim).")
-    if op == "transpose_last2":
-        return _ops.transpose_last2(_eval_value(args[0], ctx, params, u_elem=u_elem))
-    if op == "einsum":
-        subscripts = args[0]
-        operands = [_eval_value(arg, ctx, params, u_elem=u_elem) for arg in args[1:]]
-        return jnp.einsum(subscripts, *operands)
-
-    raise ValueError(f"Unknown Expr op: {op}")
+def _eval_expr(
+    expr: Expr,
+    ctx: VolumeContext | SurfaceContext,
+    params: ParamsLike,
+    u_elem: UElement | None = None,
+):
+    plan = make_eval_plan(expr)
+    return eval_with_plan(plan, ctx, params, u_elem=u_elem)
 
 
 __all__ = [
@@ -803,6 +1240,7 @@ __all__ = [
     "compile_mixed_residual",
     "grad",
     "sym_grad",
+    "outer",
     "dot",
     "ddot",
     "inner",
@@ -814,5 +1252,7 @@ __all__ = [
     "transpose",
     "log",
     "transpose_last2",
+    "matmul",
+    "matmul_std",
     "einsum",
 ]