Differentiation¶

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 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 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.
  """
  rows, inverse = np.unique(np.asarray(sparsity.rows, dtype=np.int64), return_inverse=True)
  mask = _incidence((len(rows), sparsity.shape[1]), inverse, np.asarray(sparsity.cols, dtype=np.int64))
  csr, csc = mask.tocsr(), mask.tocsc()
  row_ptr, col_ind = csr.indptr, csr.indices
  col_ptr, row_ind = csc.indptr, csc.indices
  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) -> Expr:
  """Sparse Jacobian: a sparse tensor of shape ``(expr.size, wrt.size)`` storing only its structural nonzeros.

  ``wrt`` must be an independent input from ``sc.sym`` or a Function declaration.
  Takes the structured path when ``expr`` is a mapped loop (or a concatenation of them) over
  ``wrt`` or a constant gather of it. It colors 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.
  """
  return _sparse_jacobian(expr, wrt).leaf()

def sparse_jacobian_reference(expr: Expr, wrt: Expr) -> Expr:
  """Reference sparse Jacobian path: build the dense Jacobian and keep its structural 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 _Compact(sparsity, gather(dense, flat)).leaf()