Building functions¶

def group(*parts: Tree[Any, Any]) -> Tree[Any, Any]:
  """Group one to eight trees side by side. Nest groups for greater widths."""
  return _G(parts)

def function(
  *inputs: Tree[Any, Any], outputs: Tree[Any, Any] | None = None, name: str | None = None
) -> Callable[[Callable[..., Any]], Function[Any, Any, Any, Any]]:
  """Trace a callable with one declaration per parameter and an optional output tree."""

  def decorate(fn: Callable[..., Any]) -> Function[Any, Any, Any, Any]:
    function_name = name or getattr(fn, "__name__", "fn")
    parameters = tuple(inspect.signature(fn).parameters.values())
    if (inputs or outputs is not None) and (
      len(parameters) != len(inputs) or any(p.kind not in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) for p in parameters)
    ):
      raise TypeError(f"{function_name}: declare one tree per positional parameter, got {len(inputs)} declarations for {len(parameters)} parameters")
    return _traced(Function(function_name, fn, None if not inputs and outputs is None and parameters else parameter_list(inputs), outputs))

  return decorate

def jacobian(
  source: Expr | Function | ConcreteFunction,
  *args: Expr | str,
  wrt: Expr | str | None = None,
  of: str | None = None,
  name: str | None = None,
) -> Expr | Function:
  """Build a dense Jacobian for an Expr or a named Function output.

  For an expression source, ``wrt`` must be an independent input.
  """
  if isinstance(source, Expr):
    return _jacobian_expr(source, _expr_wrt("jacobian", args, wrt, of, name))
  if not isinstance(source, Function | ConcreteFunction):
    raise TypeError("jacobian source must be an Expr or Function")
  function_of, function_wrt = _function_names("jacobian", args, wrt, of)
  return _unseeded(source, name, function_of, function_wrt, Jac)

def gradient(
  source: Expr | Function | ConcreteFunction,
  *args: Expr | str,
  wrt: Expr | str | None = None,
  of: str | None = None,
  name: str | None = None,
) -> Expr | Function:
  """Build a gradient for an Expr or a named Function output.

  For an expression source, ``wrt`` must be an independent input.
  """
  if isinstance(source, Expr):
    return _gradient_expr(source, _expr_wrt("gradient", args, wrt, of, name))
  if not isinstance(source, Function | ConcreteFunction):
    raise TypeError("gradient source must be an Expr or Function")
  function_of, function_wrt = _function_names("gradient", args, wrt, of)
  return _unseeded(source, name, function_of, function_wrt, Grad)

def hessian(
  source: Expr | Function | ConcreteFunction,
  *args: Expr | str,
  wrt: Expr | str | None = None,
  of: str | None = None,
  name: str | None = None,
) -> Expr | Function:
  """Build a Hessian for an Expr or a named Function output.

  For an expression source, ``wrt`` must be an independent input.
  """
  if isinstance(source, Expr):
    return _hessian_expr(source, _expr_wrt("hessian", args, wrt, of, name))
  if not isinstance(source, Function | ConcreteFunction):
    raise TypeError("hessian source must be an Expr or Function")
  function_of, function_wrt = _function_names("hessian", args, wrt, of)
  return _unseeded(source, name, function_of, function_wrt, Hess)

def sparse_jacobian(
  source: Expr | Function | ConcreteFunction,
  *args: Expr | str,
  wrt: Expr | str | None = None,
  of: str | None = None,
  name: str | None = None,
  layout: Layout = "csc",
) -> Expr | Function:
  """Build a sparse Jacobian that stores only its structural nonzeros, for an Expr or a named Function output.

  For an expression source, ``wrt`` must be an independent input, and the result is a sparse
  tensor of shape ``(source.size, wrt.size)``. For a Function, the output is a sparse leaf whose
  values are stored in ``layout``: ``"csc"`` (the default), ``"csr"``, or a ``SparsityPattern``
  listing the same coordinates in the order the caller wants.
  """
  if isinstance(source, Expr):
    if layout != "csc":
      raise TypeError("sparse_jacobian Expr form does not accept layout")
    return _expr_sparse_jacobian(source, _expr_wrt("sparse_jacobian", args, wrt, of, name))
  if not isinstance(source, Function | ConcreteFunction):
    raise TypeError("sparse_jacobian source must be an Expr or Function")
  function_of, function_wrt = _function_names("sparse_jacobian", args, wrt, of)
  return _unseeded(source, name, function_of, function_wrt, SpJac, layout=layout)

def sparse_hessian(
  source: Expr | Function | ConcreteFunction,
  *args: Expr | str,
  wrt: Expr | str | None = None,
  of: str | None = None,
  name: str | None = None,
  triangle: Triangle = "full",
  layout: Layout = "csc",
) -> Expr | Function:
  """Build a sparse Hessian that stores only its structural nonzeros, for an Expr or a named scalar Function output.

  For an expression source, ``wrt`` must be an independent input, and the result is a sparse
  tensor of shape ``(wrt.size, wrt.size)``. ``triangle`` selects the stored entries: ``"full"``
  keeps the whole symmetric pattern, while ``"lower"`` and ``"upper"`` keep one triangle with the
  diagonal. For a Function, the output is a sparse leaf whose values are stored in ``layout``, as
  for ``sparse_jacobian``.
  """
  if isinstance(source, Expr):
    if layout != "csc":
      raise TypeError("sparse_hessian Expr form does not accept layout")
    return _expr_sparse_hessian(source, _expr_wrt("sparse_hessian", args, wrt, of, name), triangle=triangle)
  if not isinstance(source, Function | ConcreteFunction):
    raise TypeError("sparse_hessian source must be an Expr or Function")
  function_of, function_wrt = _function_names("sparse_hessian", args, wrt, of)
  return _unseeded(source, name, function_of, function_wrt, SpHess, triangle=triangle, layout=layout)

def forward(fn: Function | ConcreteFunction, of: str | None = None, wrt: str | None = None, *, name: str | None = None) -> Function:
  """Create a forward derivative with one seed parameter shaped and typed as ``wrt``."""
  return _seeded(fn, of, wrt, name, "fwd")

def adjoint(fn: Function | ConcreteFunction, of: str | None = None, wrt: str | None = None, *, name: str | None = None) -> Function:
  """Create an adjoint derivative with one seed parameter shaped and typed as ``of``."""
  return _seeded(fn, of, wrt, name, "adj")

def lagrangian_hessian[*Ss, *Ns, SO, NO](
  fn: Function[tuple[*Ss], tuple[*Ns], SO, NO], wrt: str | None = None, *, name: str | None = None, aux_name: str = "gamma"
) -> Function[tuple[*Ss, SO], tuple[*Ns, NO], Expr, Array]:
  """Create the dense weighted Hessian with one multiplier parameter matching the output tree."""
  return _seeded(fn, None, wrt, name, "hess", aux_name)

def sparse_lagrangian_hessian[*Ss, *Ns, SO, NO](
  fn: Function[tuple[*Ss], tuple[*Ns], SO, NO],
  wrt: str | None = None,
  *,
  name: str | None = None,
  aux_name: str = "gamma",
  triangle: Triangle = "full",
  layout: Layout = "csc",
) -> Function[tuple[*Ss, SO], tuple[*Ns, NO], Expr, NumericalArray]:
  """Create the sparse weighted Hessian with multipliers matching the output tree, stored in ``layout``."""
  return _seeded(fn, None, wrt, name, "sphess", aux_name, triangle, layout)

def custom_derivative[SI, NI, SO, NO](
  fn: Function[SI, NI, SO, NO] | ConcreteFunction[SI, NI, SO, NO],
  *,
  jvp: _Rule | None = None,
  fwd: _Rule | None = None,
  bwd: _Rule | None = None,
  sparsity: Callable[[str, str, tuple[int, int]], SparsityPattern | np.ndarray] | None = None,
) -> Function[SI, NI, SO, NO]:
  """Return ``fn`` with derivative rules that replace differentiating its body.

  ``jvp(*inputs, *directions)`` is the forward rule, the product of the Jacobian with a direction.
  It takes ``fn``'s parameters and then one direction per parameter, and returns the change of
  each output along that direction, in ``fn``'s output structure. Several directions apply the
  rule to each direction in one mapped loop, and a direction the rule never reads is never computed.

  ``fwd`` and ``bwd`` form the reverse rule, the product of output weights with the Jacobian.
  ``fwd(*inputs)`` returns ``(outputs, residuals)``, and the outputs become ``fn``'s.
  ``bwd(residuals, weights)`` takes one weight per output and returns one weighted derivative per
  parameter of ``fn``. A call computes its outputs and residuals together, so the reverse rule
  reuses them instead of computing them again, and a second derivative still differentiates them.

  A mode without a rule, forward or reverse, differentiates the body. Rules may be shape templates,
  which take the shapes of each call of ``fn``. ``sparsity(of, wrt, shape)`` gives the Jacobian pattern of output
  ``of`` in input ``wrt``, as a ``SparsityPattern`` or a boolean mask of ``shape``, which is
  ``(of.size, wrt.size)``. Without it, the pattern of a function with rules is dense.
  """
  if (fwd is None) != (bwd is None):
    raise TypeError("custom_derivative takes fwd and bwd together")
  if jvp is None and fwd is None and sparsity is None:
    raise TypeError("custom_derivative needs at least one of jvp, fwd and bwd, or sparsity")
  if not isinstance(fn, Function | ConcreteFunction):
    raise TypeError("sc.custom_derivative needs a Function; wrap an extern call in a Function first")
  source = Function._from_instance(fn) if isinstance(fn, ConcreteFunction) else fn
  return lift(source, source.inputs, source.outputs, lambda concrete: _with_rules(concrete, jvp, fwd, bwd, sparsity), name=source.name)

def broadcast[T](value: T, /) -> T:
  """Mark one tensor argument to be passed whole to every mapped iteration."""
  return cast(T, _Broadcast(value))

def window[T](value: T, start: int, stride: int, /) -> T:
  """Mark rank-1 windows whose width comes from the callee's declared shape.

  Iteration ``i`` reads at ``start + i * stride``. For a shape hole, use a leading-axis
  view or explicitly instantiate the callee to specify the window width.
  """
  start, stride = int(start), int(stride)
  if start < 0 or stride < 0:
    raise ValueError("window start and stride must be non-negative")
  return cast(T, _Window(value, start, stride))