Building functions¶

class Tree[Symbolic, Numerical]:
  """A pytree declaration whose leaves are ``Expr`` symbolically and NumPy arrays numerically.

  A tree *is* its structure: only ``G`` introduces a tuple, so a one-leaf tree is the bare leaf
  and never a one-element tuple. See ``L`` for what that means at a call site.
  """

  names: tuple[str, ...]
  decls: tuple[LeafDecl, ...]

  @property
  def shapes(self) -> tuple[tuple[int, ...], ...]:
    """The leaf shapes in C-signature order."""
    if any(decl is Ellipsis for decl in self.decls):
      raise TypeError(f"tree {self.names} has inferred shapes; resolve them by tracing first")
    return tuple(cast(TensorType, decl).shape for decl in self.decls)

  @property
  def types(self) -> tuple[TensorType, ...]:
    """The leaf tensor types in C-signature order."""
    if any(decl is Ellipsis for decl in self.decls):
      raise TypeError(f"tree {self.names} has inferred shapes; resolve them by tracing first")
    return cast(tuple[TensorType, ...], self.decls)

  @property
  def size(self) -> int:
    """The number of leaves."""
    return len(self.names)

  def symbols(self, *, diff: bool | None = None) -> Symbolic:
    """Create named input expressions with this tree's structure."""
    raise NotImplementedError

  def relabel(self, prefix: str) -> Tree[Symbolic, Numerical]:
    """Return the same structure with ``prefix`` added to every leaf name."""
    raise NotImplementedError

  def with_types(self, types: tuple[TensorType, ...]) -> Tree[Symbolic, Numerical]:
    """Return the same structure with resolved leaf types."""
    raise NotImplementedError

  def resolved(self, traced: tuple[TensorType, ...]) -> tuple[TensorType, ...]:
    """Check traced output types against this declaration and resolve inferred shapes."""
    if len(traced) != self.size:
      raise TypeError(f"declared {self.size} outputs {self.names}, body returned {len(traced)}")
    for name, decl, actual in zip(self.names, self.decls, traced, strict=True):
      if decl is not Ellipsis and (decl.shape != actual.shape or decl.dtype != actual.dtype):
        raise TypeError(f"{name!r} declared with type {decl}, traced type {actual}")
    return traced

  def index(self, name: str) -> int:
    """Return the flat index for ``name``, or raise with the declared choices."""
    if name not in self.names:
      raise ValueError(f"unknown name {name!r}; declared {self.names}")
    return self.names.index(name)

  def is_symbolic(self, value: Symbolic | Numerical, /) -> TypeGuard[Symbolic]:
    """Whether ``value`` has at least one leaf and every leaf is an ``Expr``.

    This is the leaf-kind half of ``Function.__call__``'s dispatch. It deliberately ignores
    structure so that a wrongly-shaped tree is reported by ``flatten_symbolic`` against the
    declared names instead of being rejected here as a kind mismatch.
    """
    leaves = _leaves(value)
    return bool(leaves) and all(isinstance(leaf, Expr) for leaf in leaves)

  def is_numerical(self, value: Symbolic | Numerical, /) -> TypeGuard[Numerical]:
    """Whether no leaf of ``value`` is an ``Expr``.

    The numerical side is the fallback: array-likes are coerced by ``flatten_numerical``, so a
    leaf only has to *not* be symbolic. Values with no leaves land here and are reported as a
    structure error rather than as a mixed call.
    """
    return not any(isinstance(leaf, Expr) for leaf in _leaves(value))

  def flatten_symbolic(self, value: Symbolic, what: str, *, allow_scalar: bool = False) -> tuple[Expr, ...]:
    """Validate and flatten a symbolic value."""
    raise NotImplementedError

  def flatten_numerical(self, value: Numerical, what: str) -> tuple[np.ndarray, ...]:
    """Validate and flatten a numerical value."""
    raise NotImplementedError

  def unflatten(self, values: tuple[Any, ...]) -> Any:
    """Rebuild this tree's structure from flat values."""
    raise NotImplementedError

  def _check_unique(self) -> None:
    if len(set(self.names)) != len(self.names):
      raise ValueError(f"duplicate names in {self.names}")

class L(Tree[Expr, np.ndarray]):
  """Declare one named tensor.

  The declared name is external metadata and need not match the local name used by a decorated
  function body. Pass a ``TensorType`` to set dtype or differentiability explicitly.

  **A single leaf is passed and returned unpacked.** An ``L`` input tree takes the tensor itself,
  not ``(tensor,)``, and an ``L`` output tree returns the tensor itself, not a one-element tuple.
  Do not destructure a single-leaf result: ``(y,) = fn(x)`` does not raise, it iterates the
  returned tensor along its first axis exactly as NumPy would.
  """

  def __init__(self, name: str, shape: ShapeDecl, /) -> None:
    if not isinstance(name, str) or not name:
      raise ValueError("L needs a non-empty name")
    self.names = (name,)
    if shape is Ellipsis:
      self.decls = (Ellipsis,)
    elif isinstance(shape, TensorType):
      self.decls = (shape,)
    else:
      self.decls = (TensorType(as_shape(shape)),)

  def symbols(self, *, diff: bool | None = None) -> Expr:
    type_ = self.types[0]
    if diff is not None:
      type_ = TensorType(type_.shape, type_.dtype, diff)
    return Expr(ExprOp.INPUT, type=type_, name=self.names[0])

  def relabel(self, prefix: str) -> L:
    return L(prefix + self.names[0], self.decls[0])

  def with_types(self, types: tuple[TensorType, ...]) -> L:
    if len(types) != 1:
      raise ValueError(f"L expects one resolved type, got {len(types)}")
    return L(self.names[0], types[0])

  def flatten_symbolic(self, value: Expr, what: str, *, allow_scalar: bool = False) -> tuple[Expr, ...]:
    if not isinstance(value, Expr):
      raise ValueError(f"{what}: expected an Expr for {self.names[0]!r}, got {type(value).__name__}")
    decl = self.decls[0]
    if decl is not Ellipsis and value.shape != decl.shape and not (allow_scalar and value.shape == ()):
      raise ValueError(f"{what}: expected shape {decl.shape} for {self.names[0]!r}, got {value.shape}")
    return (value,)

  def flatten_numerical(self, value: np.ndarray, what: str) -> tuple[np.ndarray, ...]:
    if isinstance(value, Expr):
      raise ValueError(f"{what}: expected a numerical value for {self.names[0]!r}, got Expr")
    array = np.asarray(value, dtype=self.types[0].dtype.numpy())
    if array.shape != self.shapes[0]:
      raise ValueError(f"{what}: expected shape {self.shapes[0]} for {self.names[0]!r}, got {array.shape}")
    return (np.require(array, requirements="C"),)

  def unflatten(self, values: tuple[Any, ...]) -> Any:
    if len(values) != 1:
      raise ValueError(f"L expects one flat value, got {len(values)}")
    return values[0]

@dataclass(frozen=True, slots=True)
class DerivSpec:
  """Base class for typed requests passed to Function.factory."""

  kind: ClassVar[str]
  of: str
  wrt: str

  @property
  def output_name(self) -> str:
    return f"{self.kind}_{self.of}_{self.wrt}"

  def build(self, inputs: Mapping[str, Expr], outputs: Mapping[str, Expr]) -> tuple[Expr, SparsityPattern | None, int | None]:
    raise NotImplementedError

  def _in(self, inputs: Mapping[str, Expr], name: str) -> Expr:
    if name not in inputs:
      raise ValueError(f"unknown factory input {name!r} in output {self}")
    return inputs[name]

  def _out(self, outputs: Mapping[str, Expr], name: str) -> Expr:
    if name not in outputs:
      raise ValueError(f"unknown factory output {name!r} in output {self}")
    return outputs[name]

@dataclass(frozen=True, slots=True)
class Jac(DerivSpec):
  """Dense Jacobian ``d of / d wrt``."""

  kind = "jac"

  def build(self, inputs: Mapping[str, Expr], outputs: Mapping[str, Expr]) -> tuple[Expr, SparsityPattern | None, int | None]:
    return jacobian(self._out(outputs, self.of), self._in(inputs, self.wrt)), None, None

@dataclass(frozen=True, slots=True)
class Grad(DerivSpec):
  """Gradient of a scalar output."""

  kind = "grad"

  def build(self, inputs: Mapping[str, Expr], outputs: Mapping[str, Expr]) -> tuple[Expr, SparsityPattern | None, int | None]:
    return gradient(self._out(outputs, self.of), self._in(inputs, self.wrt)), None, None

@dataclass(frozen=True, slots=True)
class Hess(DerivSpec):
  """Dense Hessian of a scalar output."""

  kind = "hess"

  @property
  def output_name(self) -> str:
    return f"{self.kind}_{self.of}_{self.wrt}_{self.wrt}"

  def build(self, inputs: Mapping[str, Expr], outputs: Mapping[str, Expr]) -> tuple[Expr, SparsityPattern | None, int | None]:
    y = self._out(outputs, self.of)
    x0 = self._in(inputs, self.wrt)
    return hessian(y, x0), None, None

@dataclass(frozen=True, slots=True)
class SpJac(DerivSpec):
  """Compact nonzero Jacobian values, with the sparsity that indexes them."""

  kind = "spjac"

  def build(self, inputs: Mapping[str, Expr], outputs: Mapping[str, Expr]) -> tuple[Expr, SparsityPattern | None, int | None]:
    sj = sparse_jacobian(self._out(outputs, self.of), self._in(inputs, self.wrt))
    return sj.values, sj.sparsity, sj.coloring_width

@dataclass(frozen=True, slots=True)
class SpHess(DerivSpec):
  """Compact nonzero Hessian values, with the sparsity that indexes them."""

  kind = "sphess"
  triangle: Triangle = field(default="full", kw_only=True)

  def __post_init__(self) -> None:
    _validate_triangle(self.triangle)

  @property
  def output_name(self) -> str:
    return f"{self.kind}_{self.of}_{self.wrt}_{self.wrt}"

  def build(self, inputs: Mapping[str, Expr], outputs: Mapping[str, Expr]) -> tuple[Expr, SparsityPattern | None, int | None]:
    y = self._out(outputs, self.of)
    x0 = self._in(inputs, self.wrt)
    sh = sparse_hessian(y, x0, triangle=self.triangle)
    return sh.values, sh.sparsity, sh.coloring_width

@dataclass(frozen=True, slots=True)
class Fwd(DerivSpec):
  """Seeded forward mode: ``J(of, wrt) @ fwd:wrt``."""

  kind = "fwd"

  def build(self, inputs: Mapping[str, Expr], outputs: Mapping[str, Expr]) -> tuple[Expr, SparsityPattern | None, int | None]:
    return jvp(self._out(outputs, self.of), self._in(inputs, self.wrt), self._in(inputs, f"fwd:{self.wrt}")), None, None

@dataclass(frozen=True, slots=True)
class Adj(DerivSpec):
  """Seeded reverse mode: ``J(of, wrt).T @ lam:of``."""

  kind = "adj"

  def build(self, inputs: Mapping[str, Expr], outputs: Mapping[str, Expr]) -> tuple[Expr, SparsityPattern | None, int | None]:
    return vjp((self._out(outputs, self.of),), (self._in(inputs, self.wrt),), (self._in(inputs, f"lam:{self.of}"),))[0], None, None