Differentiation¶

def jvp(expr: Expr, wrt: Expr, seed: Expr) -> Expr:
  """Forward-mode derivative: ``J(expr, wrt) @ seed``, with ``seed`` shaped like ``wrt``.

  ``wrt`` must be an independent input from ``sc.sym`` or a Function declaration.
  """
  tangent = jvp_many(expr, wrt, seed.reshape((1, *seed.shape)))
  if tangent.op == ExprOp.CONST:
    assert tangent.value is not None
    return Expr.const(tangent.value.reshape(expr.shape), dtype=tangent.type.dtype, lowering=tangent.lowering)
  return tangent.reshape(expr.shape)

def jvp_many(expr: Expr, wrt: Expr, seeds: Expr) -> Expr:
  """Forward mode over seeds shaped ``(nseed, *wrt.shape)``, returning ``(nseed, *expr.shape)``.

  ``wrt`` must be an independent input from ``sc.sym`` or a Function declaration. One traversal
  shares primal work across all seeds. A missing derivative rule raises.
  """
  _check_wrt(wrt)
  if len(seeds.shape) < 1 or seeds.shape[1:] != wrt.shape:
    raise ValueError(f"multi-seed JVP expects seeds shape (nseed, *{wrt.shape}), got {seeds.shape}")
  if not wrt.type.dtype.is_floating or not expr.type.dtype.is_floating:
    raise TypeError("integer and boolean values have no tangent")
  if seeds.type.dtype != wrt.type.dtype:
    raise TypeError(f"JVP seeds must have primal dtype {wrt.type.dtype}, got {seeds.type.dtype}")
  seeds = simplify_cse_fixpoint(seeds)
  (ret,) = _pushforward((expr,), {wrt: seeds}, seeds.shape[0])
  assert ret is not None
  return ret

def vjp(outputs: Sequence[Expr], wrts: Sequence[Expr], cotangents: Sequence[Expr]) -> tuple[Expr, ...]:
  """Reverse-mode derivative: one adjoint per entry of ``wrts``, seeded by ``cotangents``.

  Each cotangent is shaped like its output. One sweep computes the derivative of a scalar with
  respect to every requested input at once, which is why gradients go through here.
  Every entry of ``wrts`` must be an independent input from ``sc.sym`` or a Function declaration.
  Results follow ``wrts`` order.
  """
  if len(outputs) != len(cotangents):
    raise ValueError(f"expected {len(outputs)} cotangents, got {len(cotangents)}")
  for wrt in wrts:
    _check_wrt(wrt)
  adjoints: dict[int, Expr] = {}
  nodes = topo(outputs)
  expr_ids = {e.id for e in nodes}
  dep_memo: dict[tuple[int, int], bool] = {}

  def needed(expr: Expr) -> bool:
    return any(_depends_on(expr, wrt, dep_memo, through_stops=False) for wrt in wrts)

  for out, cot in zip(outputs, cotangents, strict=True):
    if out.shape != cot.shape:
      raise ValueError(f"cotangent for output shape {out.shape} has shape {cot.shape}")
    adjoints[out.id] = cot if out.id not in adjoints else adjoints[out.id] + cot

  for expr in reversed(nodes):
    cot = adjoints.get(expr.id)
    if cot is None or expr.op in {ExprOp.INPUT, ExprOp.CONST, ExprOp.STOP_GRADIENT} or not needed(expr):
      continue
    if expr.op == ExprOp.VMAP:
      for arg, arg_cot in _vmap_vjp(expr, cot, wrts, dep_memo, pullback=vjp):
        if arg.id in expr_ids:
          adjoints[arg.id] = arg_cot if arg.id not in adjoints else adjoints[arg.id] + arg_cot
      continue
    active = tuple(i for i, arg in enumerate(expr.args) if needed(arg)) if expr.op in {ExprOp.CALL, *ELEMENTWISE} else None
    for arg, arg_cot in zip(expr.args, _local_vjp(expr, cot, active), strict=True):
      if arg_cot is not None and arg.id in expr_ids:
        adjoints[arg.id] = arg_cot if arg.id not in adjoints else adjoints[arg.id] + arg_cot

  return tuple(adjoints.get(wrt.id, zeros_like(wrt)) for wrt in wrts)

def stop_gradient(value: Any) -> Expr:
  """Return ``value`` unchanged in evaluation and treat it as a constant in derivatives.

  Forward and reverse derivatives through this expression are zero, as is its Jacobian sparsity.
  Other uses of ``value`` still contribute their derivatives. Its generated evaluation C is unchanged.
  """
  expr = as_expr(value)
  return Expr(ExprOp.STOP_GRADIENT, (expr,), expr.type, lowering=expr.lowering)

def jacobian(expr: Expr, wrt: Expr) -> Expr:
  """Dense Jacobian ``d expr / d wrt``, shape ``(expr.size, wrt.size)``.

  ``wrt`` must be an independent input from ``sc.sym`` or a Function declaration.
  Column ``j`` is the derivative with respect to ``wrt[j]``. Computed by pushing the whole
  identity through forward mode in one batched pass, then simplifying.
  """
  _check_wrt(wrt)
  if wrt.size == 0:
    return Expr.const(np.zeros((expr.size, 0), dtype=expr.type.dtype.numpy()), dtype=expr.type.dtype)
  seed_arr = np.eye(wrt.size, dtype=wrt.type.dtype.numpy()).reshape((wrt.size, *wrt.shape))
  return simplify_cse_fixpoint(jvp_many(expr, wrt, Expr.const(seed_arr, dtype=wrt.type.dtype)).reshape((wrt.size, expr.size)).transpose((1, 0)))

def gradient(expr: Expr, wrt: Expr) -> Expr:
  """Gradient of a scalar ``expr`` with respect to ``wrt``, as one reverse sweep.

  ``wrt`` must be an independent input from ``sc.sym`` or a Function declaration.
  """
  if expr.size != 1:
    raise ValueError("gradient expects a scalar expression")
  return vjp((expr,), (wrt,), (_ones_like(expr),))[0]

def hessian(expr: Expr, wrt: Expr) -> Expr:
  """Second derivatives of a scalar ``expr``: the Jacobian of its gradient.

  ``wrt`` must be an independent input from ``sc.sym`` or a Function declaration.
  """
  return jacobian(gradient(expr, wrt).reshape((wrt.size,)), wrt)

def finite_difference(fun: Any, x: np.ndarray, eps: float = 1e-6) -> np.ndarray:
  """Approximate the Jacobian of the numerical ``fun`` at ``x`` by central differences.

  Returns shape ``(output size, x.size)``. Useful as an independent check of a symbolic derivative.
  """
  x = np.asarray(x, dtype=np.float64)
  y0 = np.asarray(fun(x), dtype=np.float64).reshape(-1)
  jac = np.empty((y0.size, x.size), dtype=np.float64)
  flat = x.reshape(-1)
  for i in range(flat.size):
    xp = flat.copy()
    xm = flat.copy()
    xp[i] += eps
    xm[i] -= eps
    jac[:, i] = (np.asarray(fun(xp.reshape(x.shape))).reshape(-1) - np.asarray(fun(xm.reshape(x.shape))).reshape(-1)) / (2 * eps)
  return jac

def jacobian_sparsity(expr: Expr, wrt: Expr) -> SparsityPattern:
  """Estimate structural sparsity of ``d vec(expr) / d vec(wrt)``.

  ``wrt`` must be an independent input from ``sc.sym`` or a Function declaration.
  This is purely symbolic: it tracks element dependencies through the graph without using
  numerical values. It is conservative for nonsmooth elementwise ops and block ops, but exact
  for the structural/arithmetic subset currently implemented here.
  """
  _check_wrt(wrt)
  return _mask_sparsity(_jac_mask(expr, wrt, {}))

def column_coloring(sparsity: SparsityPattern) -> tuple[int, ...]:
  """Assign each column a color, greedily, so no two columns sharing a row get the same one.

  Columns of one color can be recovered from a single forward pass, so the number of colors is
  the number of passes a compact Jacobian costs.
  """
  row_ptr, col_ind, _ = sparsity.to_csr()
  col_ptr, row_ind, _ = sparsity.to_csc()
  colors: list[int] = []
  for col in range(sparsity.shape[1]):
    used = {colors[other] for row in row_ind[col_ptr[col] : col_ptr[col + 1]] for other in col_ind[row_ptr[row] : row_ptr[row + 1]] if other < col}
    color = 0
    while color in used:
      color += 1
    colors.append(color)
  return tuple(colors)

def star_coloring(sparsity: SparsityPattern) -> tuple[int, ...]:
  """Greedily star-color a square sparsity pattern in column order.

  The pattern is treated as an undirected graph. A valid coloring is proper, and no simple path
  of three edges has only two colors. This is the coloring needed to recover a symmetric Hessian
  from compressed forward products. It is deliberately not distance-2 coloring.
  """
  rows, cols = sparsity.shape
  if rows != cols:
    raise ValueError(f"star coloring requires a square sparsity pattern, got {sparsity.shape}")

  neighbors = [set() for _ in range(rows)]
  for row, col in zip(sparsity.rows, sparsity.cols, strict=True):
    if row != col:
      neighbors[row].add(col)
      neighbors[col].add(row)

  colors = [-1] * rows
  for vertex in range(rows):
    color = 0
    while _star_color_conflicts(vertex, color, colors, neighbors):
      color += 1
    colors[vertex] = color
  return tuple(colors)

def color_groups(colors: tuple[int, ...]) -> tuple[tuple[int, ...], ...]:
  """Invert a coloring into the column indices belonging to each color."""
  if not colors:
    return ()
  return tuple(tuple(i for i, c in enumerate(colors) if c == color) for color in range(max(colors) + 1))

def sparse_jacobian(expr: Expr, wrt: Expr) -> SparseJacobian:
  """Compact sparse Jacobian: the structural pattern plus an expression for its nonzeros only.

  ``wrt`` must be an independent input from ``sc.sym`` or a Function declaration.
  Takes the structured path when ``expr`` is a ``VMAP`` (or a concatenation of them) over exactly
  ``wrt``, coloring the callee's small local pattern instead of the whole matrix, so the cost
  tracks the callee rather than the iteration count. Otherwise falls back to coloring the global
  pattern. The value order is the pattern's ``(rows, cols)`` order, which is not necessarily
  sorted.
  """
  _check_wrt(wrt)
  structured = _sparse_jacobian_structured(expr, wrt)
  if structured is not None:
    return structured
  return sparse_jacobian_colored(expr, wrt)

def sparse_hessian(expr: Expr, wrt: Expr, *, triangle: Triangle = "full") -> SparseJacobian:
  """Compact sparse Hessian of a scalar expression using one global star-colored JVP batch.

  ``wrt`` must be an independent input from ``sc.sym`` or a Function declaration.
  """
  triangle = _validate_triangle(triangle)
  if expr.size != 1:
    raise ValueError("sparse_hessian expects a scalar expression")
  _check_wrt(wrt)
  gradient_expr = cse(simplify(gradient(expr, wrt).reshape((wrt.size,))))
  sparsity = _symmetrize_sparsity(jacobian_sparsity(gradient_expr, wrt))
  colors = star_coloring(sparsity)
  recovery = _star_recovery_indices(sparsity, colors)
  return _sparse_jacobian_colored(gradient_expr, wrt, sparsity, colors, recovery).triangle(triangle)

def sparse_jacobian_colored(expr: Expr, wrt: Expr) -> SparseJacobian:
  """Compact sparse Jacobian values from graph-colored compressed JVPs.

  ``wrt`` must be an independent input from ``sc.sym`` or a Function declaration.
  """
  _check_wrt(wrt)
  expr = cse(expr)
  sparsity = jacobian_sparsity(expr, wrt)
  colors = column_coloring(sparsity)
  return _sparse_jacobian_colored(expr, wrt, sparsity, colors)

def sparse_jacobian_reference(expr: Expr, wrt: Expr) -> SparseJacobian:
  """Reference compact Jacobian path: build dense ``J`` and gather nonzeros.

  ``wrt`` must be an independent input from ``sc.sym`` or a Function declaration.
  """

  sparsity = jacobian_sparsity(expr, wrt)
  dense = jacobian(expr, wrt)
  flat = np.asarray(sparsity.rows, dtype=np.int64) * wrt.size + np.asarray(sparsity.cols, dtype=np.int64)
  return SparseJacobian(sparsity, gather(dense, flat))