Building functions¶

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

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

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

  def __setattr__(self, name: str, value: Any) -> None:
    if name in self.__dict__:
      raise AttributeError(f"tree attribute {name!r} is read-only")
    object.__setattr__(self, name, value)

  def __delattr__(self, name: str) -> None:
    raise AttributeError(f"tree attribute {name!r} is read-only")

  @property
  def has_holes(self) -> bool:
    """Whether any leaf shape needs binding."""
    return any(decl is None for decl in self.decls)

  @property
  def shapes(self) -> tuple[tuple[int, ...], ...]:
    """The leaf shapes in C-signature order."""
    if any(decl is None 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 None 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 None 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[Array, ...]:
    """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}")

@dataclass(frozen=True, slots=True)
class DerivSpec:
  """Base class for typed requests passed to ConcreteFunction.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