Core¶

class ExprOp(StrEnum):
  """The expression dialect's operation set.

  A ``StrEnum``, so an op is a proper enum value and still prints and serializes as its name.
  ``OP_INFO`` carries the arity and the differentiability of each.
  """

  INPUT = "input"
  CONST = "const"
  NEG = "neg"
  SIN = "sin"
  COS = "cos"
  TAN = "tan"
  ASIN = "asin"
  ACOS = "acos"
  ATAN = "atan"
  SINH = "sinh"
  COSH = "cosh"
  TANH = "tanh"
  ERF = "erf"
  EXP = "exp"
  LOG = "log"
  SQRT = "sqrt"
  ABS = "abs"
  FLOOR = "floor"
  CEIL = "ceil"
  ADD = "add"
  SUB = "sub"
  MUL = "mul"
  DIV = "div"
  POW = "pow"
  ATAN2 = "atan2"
  MINIMUM = "minimum"
  MAXIMUM = "maximum"
  SUM = "sum"
  RESHAPE = "reshape"
  TRANSPOSE = "transpose"
  SLICE = "slice"
  GATHER = "gather"
  SCATTER = "scatter"
  STACK = "stack"
  CONCAT = "concat"
  MATMUL = "matmul"
  CALL = "call"
  VMAP = "vmap"
  SOLVER_CALL = "solver_call"
  PRINT = "print"
  STOP_GRADIENT = "stop_gradient"

@dataclass(frozen=True, slots=True)
class OpInfo:
  """What the compiler knows about one ``ExprOp``: its arity (``None`` for variadic) and whether it
  has derivative rules. ``OP_INFO`` holds one per op."""

  op: ExprOp
  arity: int | None
  differentiable: bool = True

  @property
  def name(self) -> str:
    return self.op.value

@dataclass(frozen=True, slots=True)
class TensorType:
  """The shape and dtype of an expression, and whether it is differentiable (``diff``)."""

  shape: tuple[int, ...] = ()
  dtype: DType = dtypes.float64
  diff: bool = True

  def __post_init__(self) -> None:
    _check_shape("tensor", self.shape)
    if not isinstance(self.dtype, DType):
      object.__setattr__(self, "dtype", as_dtype(self.dtype))

  @property
  def ndim(self) -> int:
    return len(self.shape)

  @property
  def size(self) -> int:
    return reduce(mul, self.shape, 1)

  @property
  def is_scalar(self) -> bool:
    return self.shape == () or self.shape == (1,) or self.size == 1

@dataclass(frozen=True, slots=True)
class DType:
  """Interned dtype descriptor.

  See ``dtypes`` for the canonical instances. A ``DType`` compares equal to its string ``name``,
  so ``dtypes.float64 == "float64"`` holds.
  """

  name: str
  bits: int
  c_type: str
  is_floating: bool = False
  is_integer: bool = False
  is_bool: bool = False
  is_unsigned: bool = False

  @property
  def itemsize(self) -> int:
    return self.bits // 8

  def numpy(self) -> np.dtype:
    return np.dtype(self.name)

  def __str__(self) -> str:
    return self.name

  def __eq__(self, other: object) -> bool:  # backward compat with string dtype
    if isinstance(other, DType):
      return self.name == other.name
    if isinstance(other, str):
      return self.name == other
    return NotImplemented

  def __hash__(self) -> int:
    return hash(self.name)

@dataclass(frozen=True, slots=True)
class SparsityPattern:
  """The structurally nonzero entries of a matrix, as coordinate lists ``rows`` and ``cols``.

  The entry order is the order of the values stored alongside the pattern. It is not necessarily
  row- or column-major. ``to_csr`` and ``to_csc`` return the permutation into either order.
  """

  shape: tuple[int, int]
  rows: tuple[int, ...]
  cols: tuple[int, ...]

  def __post_init__(self) -> None:
    if len(self.shape) != 2:
      raise ValueError(f"sparsity shape must be rank-2, got {self.shape}")
    _check_shape("sparsity", self.shape)
    if len(self.rows) != len(self.cols):
      raise ValueError("sparsity rows and cols must have the same length")
    if any(r < 0 or r >= self.shape[0] for r in self.rows) or any(c < 0 or c >= self.shape[1] for c in self.cols):
      raise ValueError(f"sparsity indices out of bounds for shape {self.shape}")
    if len(set(zip(self.rows, self.cols))) != len(self.rows):
      raise ValueError("sparsity indices must be unique")

  @property
  def nnz(self) -> int:
    return len(self.rows)

  @staticmethod
  def empty(shape: tuple[int, int]) -> SparsityPattern:
    return SparsityPattern(shape, (), ())

  @staticmethod
  def dense(shape: tuple[int, int]) -> SparsityPattern:
    rows, cols = np.nonzero(np.ones(shape, dtype=bool))
    return SparsityPattern(shape, tuple(int(x) for x in rows), tuple(int(x) for x in cols))

  @staticmethod
  def from_mask(mask: np.ndarray) -> SparsityPattern:
    mask = np.asarray(mask, dtype=bool)
    if mask.ndim != 2:
      raise ValueError(f"sparsity mask must be rank-2, got {mask.shape}")
    rows, cols = np.nonzero(mask)
    shape = (int(mask.shape[0]), int(mask.shape[1]))
    return SparsityPattern(shape, tuple(int(x) for x in rows), tuple(int(x) for x in cols))

  @staticmethod
  def from_csr(shape: tuple[int, int], row_ptr: Sequence[int], col_ind: Sequence[int]) -> SparsityPattern:
    row_ptr = tuple(int(x) for x in row_ptr)
    col_ind = tuple(int(x) for x in col_ind)
    _check_compressed_ptr("row_ptr", row_ptr, shape[0], len(col_ind))
    if any(c < 0 or c >= shape[1] for c in col_ind):
      raise ValueError(f"CSR column indices out of bounds for shape {shape}")
    rows = tuple(r for r in range(shape[0]) for _ in range(row_ptr[r + 1] - row_ptr[r]))
    return SparsityPattern(shape, rows, col_ind)

  @staticmethod
  def from_csc(shape: tuple[int, int], col_ptr: Sequence[int], row_ind: Sequence[int]) -> SparsityPattern:
    col_ptr = tuple(int(x) for x in col_ptr)
    row_ind = tuple(int(x) for x in row_ind)
    _check_compressed_ptr("col_ptr", col_ptr, shape[1], len(row_ind))
    if any(r < 0 or r >= shape[0] for r in row_ind):
      raise ValueError(f"CSC row indices out of bounds for shape {shape}")
    cols = tuple(c for c in range(shape[1]) for _ in range(col_ptr[c + 1] - col_ptr[c]))
    return SparsityPattern(shape, row_ind, cols)

  def to_mask(self) -> np.ndarray:
    mask = np.zeros(self.shape, dtype=bool)
    mask[list(self.rows), list(self.cols)] = True
    return mask

  def to_csr(self) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
    """``(row_ptr, col_ind, val_perm)``: ``val_perm[k]`` is the COO position of CSR slot ``k``, so
    ``values_csr[k] = values[val_perm[k]]`` pairs a compact COO-ordered value buffer with the CSR
    indices. The COO order is arbitrary, for example piece by piece for a mapped sparse Jacobian."""
    order = sorted(range(self.nnz), key=lambda i: (self.rows[i], self.cols[i]))
    row_ptr = [0] * (self.shape[0] + 1)
    for i in order:
      row_ptr[self.rows[i] + 1] += 1
    for r in range(self.shape[0]):
      row_ptr[r + 1] += row_ptr[r]
    return tuple(row_ptr), tuple(self.cols[i] for i in order), tuple(order)

  def to_csc(self) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
    """``(col_ptr, row_ind, val_perm)``, the compressed-column counterpart of ``to_csr``."""
    order = sorted(range(self.nnz), key=lambda i: (self.cols[i], self.rows[i]))
    col_ptr = [0] * (self.shape[1] + 1)
    for i in order:
      col_ptr[self.cols[i] + 1] += 1
    for c in range(self.shape[1]):
      col_ptr[c + 1] += col_ptr[c]
    return tuple(col_ptr), tuple(self.rows[i] for i in order), tuple(order)

class Spec(Generic[Node, Op]):
  """A verifier: a set of ``Rule``s, indexed by op so checking a node touches only its own.

  Rules with ``op=None`` apply to every node. Both dialects build their specs from this class.
  """

  def __init__(self, rules: Iterable[Rule[Node, Op]]) -> None:
    self.any: list[Rule[Node, Op]] = []
    self.by_op: dict[Op, list[Rule[Node, Op]]] = defaultdict(list)
    for rule in rules:
      if rule.op is None:
        self.any.append(rule)
      else:
        self.by_op[rule.op].append(rule)

  def candidates(self, op: Op) -> Iterable[Rule[Node, Op]]:
    yield from self.any
    yield from self.by_op.get(op, ())

  def check(self, node: Node) -> tuple[Rule[Node, Op], str] | None:
    for rule in self.candidates(getattr(node, "op")):
      diag = rule.check(node)
      if diag is not None:
        return rule, diag
    return None

  def merge(self, *others: Spec[Node, Op]) -> Spec[Node, Op]:
    rules = list(self.any)
    for op_rules in self.by_op.values():
      rules.extend(op_rules)
    for other in others:
      rules.extend(other.any)
      for op_rules in other.by_op.values():
        rules.extend(op_rules)
    return Spec(rules)

@dataclass(frozen=True, slots=True)
class Rule(Generic[Node, Op]):
  """One verifier check: ``check`` returns an error message for a bad node, or ``None``.

  The rule applies to nodes whose op is ``op``, or to every node when ``op`` is ``None``.
  """

  op: Op | None
  description: str
  check: Callable[[Node], str | None]

  def applies(self, node: Node) -> bool:
    return self.op is None or getattr(node, "op") == self.op

class VerifyError(Exception):
  """Raised by ``verify_expr`` and ``verify_program`` at the first node failing a rule.

  The message names the node, its op and the rule it broke. Both dialects raise this one type.
  """

@dataclass(frozen=True, slots=True)
class Pattern[Node: _HasOpArgs]:
  """One rewrite: nodes with op ``op`` (any op when ``None``) that satisfy ``predicate`` become ``replacement(node)``."""

  op: Hashable | None
  predicate: Callable[[Node], bool]
  replacement: Callable[[Node], Node]

  def matches(self, node: Node) -> bool:
    return (self.op is None or node.op == self.op) and self.predicate(node)

class PatternMatcher[Node: _HasOpArgs]:
  """An op-indexed set of rewrite ``Pattern``s.

  Indexing by op means a node is only tested against patterns that could match it, which is what
  keeps rewriting cheap as the pattern set grows.
  """

  def __init__(self, patterns: Iterable[Pattern[Node]]):
    self.any: list[Pattern[Node]] = []
    self.by_op: dict[Hashable, list[Pattern[Node]]] = defaultdict(list)
    for pattern in patterns:
      if pattern.op is None:
        self.any.append(pattern)
      else:
        self.by_op[pattern.op].append(pattern)

  def candidates(self, op: Hashable) -> Iterable[Pattern[Node]]:
    yield from self.any
    yield from self.by_op.get(op, ())

  def rewrite(self, node: Node) -> Node | None:
    for pattern in self.candidates(node.op):
      if pattern.matches(node):
        ret = pattern.replacement(node)
        if ret is not node:
          return ret
    return None