Core¶

def as_dtype(value: DType | str | None) -> DType:
  """Coerce a dtype name, a ``DType`` or ``None`` into a canonical ``DType``.

  ``None`` means ``float64``, scaly's default.
  """
  if value is None:
    return dtypes.float64
  if isinstance(value, DType):
    return value
  if isinstance(value, str):
    return dtypes.from_name(value)
  raise TypeError(f"cannot interpret {value!r} as a DType")

def dot(x: Any, y: Any) -> Expr:
  """Scalar product of two expressions with the same number of entries, whatever their shapes."""
  x, y = as_expr(x), as_expr(y)
  if x.size != y.size:
    raise ValueError(f"dot size mismatch: {x.shape} has {x.size} entries, {y.shape} has {y.size}")
  return (x.vec() * y.vec()).sum()

def sumsqr(x: Any) -> Expr:
  """Sum of squares of every entry: ``dot(x, x)``, a scalar."""
  x = as_expr(x)
  return dot(x, x)

def norm_2(x: Any) -> Expr:
  """Euclidean norm of every entry, as a scalar."""
  return sumsqr(x).sqrt()

def split(x: Any, sections: int | Iterable[int], *, axis: int = 0) -> tuple[Expr, ...]:
  """Split ``x`` along ``axis``, into ``sections`` equal parts or at the given boundaries."""
  x = as_expr(x)
  axis = axis if axis >= 0 else axis + len(x.shape)
  if axis < 0 or axis >= len(x.shape):
    raise ValueError(f"split axis {axis} out of bounds for shape {x.shape}")
  if isinstance(sections, int):
    if sections <= 0 or x.shape[axis] % sections != 0:
      raise ValueError(f"cannot split axis of length {x.shape[axis]} into {sections} equal sections")
    sizes = (x.shape[axis] // sections,) * sections
  else:
    sizes = tuple(int(s) for s in sections)
    if any(s < 0 for s in sizes):
      raise ValueError(f"split sizes {sizes} cannot contain negative entries")
    if sum(sizes) != x.shape[axis]:
      raise ValueError(f"split sizes {sizes} do not sum to axis length {x.shape[axis]}")
  ret: list[Expr] = []
  start = 0
  for size in sizes:
    index = [slice(None)] * len(x.shape)
    index[axis] = slice(start, start + size)
    ret.append(x[tuple(index)])
    start += size
  return tuple(ret)

def vec(x: Any) -> Expr:
  """Flatten ``x`` to rank 1 in row-major order."""
  return as_expr(x).vec()

def atan2(y: Any, x: Any) -> Expr:
  """Two-argument arctangent, elementwise: the angle of the point ``(x, y)``."""
  return binary(ExprOp.ATAN2, as_expr(y), as_expr(x))

def minimum(x: Any, y: Any) -> Expr:
  """Elementwise minimum. Non-smooth, so the result is marked non-differentiable."""
  return binary(ExprOp.MINIMUM, as_expr(x), as_expr(y))

def maximum(x: Any, y: Any) -> Expr:
  """Elementwise maximum. Non-smooth, so the result is marked non-differentiable."""
  return binary(ExprOp.MAXIMUM, as_expr(x), as_expr(y))

def load_torch_state_dict(path: str | Path) -> dict[str, np.ndarray]:
  """Load a simple PyTorch ``state_dict`` zip checkpoint without importing torch.

  Supports the modern ``torch.save(state_dict, ...)`` zip layout: ``data.pkl``
  stores tensor metadata and ``data/<n>`` stores raw CPU storage bytes. This is
  intentionally narrow. Unsupported pickle globals fail loudly instead of being
  materialized.
  """
  p = Path(path)
  storage_blobs: dict[str, bytes] = {}

  def rebuild_tensor(storage: tuple[Any, ...], storage_offset: int, size: tuple[int, ...], stride: tuple[int, ...], *args: Any) -> np.ndarray:
    del args
    if len(storage) < 5 or storage[0] != "storage":
      raise TypeError(f"unsupported torch storage persistent id: {storage!r}")
    dtype = _torch_dtype(storage[1])
    key, numel = str(storage[2]), int(storage[4])
    raw = storage_blobs[key]
    base = np.frombuffer(raw, dtype=dtype, count=numel)
    shape = tuple(int(x) for x in size)
    strides = tuple(int(x) * dtype.itemsize for x in stride)
    if not shape:
      return np.asarray(base[int(storage_offset)], dtype=dtype)
    return np.lib.stride_tricks.as_strided(base[int(storage_offset) :], shape=shape, strides=strides).copy()

  def rebuild_parameter(data: np.ndarray, *args: Any) -> np.ndarray:
    del args
    return data

  class Parameter:
    tensor: np.ndarray

    def __setstate__(self, state: tuple[np.ndarray, ...]) -> None:
      self.tensor = state[0]

  class TorchPickle(pickle.Unpickler):
    def find_class(self, module: str, name: str) -> Any:
      if module == "torch" and name in _TORCH_STORAGE_DTYPES:
        return _TORCH_STORAGE_DTYPES[name]
      if module == "torch" and name in _TORCH_DTYPES:
        return _TORCH_DTYPES[name]
      if module in {"torch._utils", "torch._tensor"} and name in {"_rebuild_tensor", "_rebuild_tensor_v2"}:
        return rebuild_tensor
      if module in {"torch._utils", "torch.nn.parameter"} and name in {"_rebuild_parameter", "_rebuild_parameter_with_state"}:
        return rebuild_parameter
      if module == "torch.nn.parameter" and name == "Parameter":
        return Parameter
      if module in {"collections", "numpy", "numpy.core.multiarray", "_codecs", "builtins"}:
        return super().find_class(module, name)
      raise pickle.UnpicklingError(f"unsupported global in torch checkpoint: {module}.{name}")

    def persistent_load(self, pid: Any) -> Any:
      return pid

  if not zipfile.is_zipfile(p):
    raise ValueError(f"unsupported legacy non-zip PyTorch checkpoint: {p}")
  with zipfile.ZipFile(p) as zf:
    names = zf.namelist()
    base = names[0].split("/", 1)[0]
    storage_blobs = {name.rsplit("/", 1)[-1]: zf.read(name) for name in names if name.startswith(f"{base}/data/")}
    with zf.open(f"{base}/data.pkl") as data:
      state = TorchPickle(data).load()
  if not isinstance(state, dict):
    raise TypeError(f"expected a state_dict in {p}, got {type(state).__name__}")
  return {k: v.tensor if isinstance(v, Parameter) else v for k, v in state.items()}

def verify_expr(root: Expr | Iterable[Expr], spec: "Spec | None" = None) -> None:
  """Topologically walk the DAG below ``root`` and raise on the first violation.

  ``spec`` defaults to ``spec_expr``, the full expression-dialect contract.
  """
  if spec is None:
    spec = spec_expr
  outputs = (root,) if isinstance(root, Expr) else tuple(root)
  for node in topo(outputs):
    result = spec.check(node)
    if result is not None:
      rule, diag = result
      label = node.name or f"%{node.id}"
      raise VerifyError(f"verify_expr: node {label} op={ExprOp(node.op).value} failed rule {rule.description!r}: {diag}")

def rewrite[Node: _HasOpArgs](
  root: Node,
  patterns: Iterable[Pattern[Node]] | PatternMatcher[Node],
  *,
  rebuild: Callable[[Node, tuple[Node, ...]], Node] | None = None,
  fixpoint: bool = True,
  revisit: bool = False,
  max_steps: int = 1_000_000,
  memo: dict[Node, Node] | None = None,
) -> Node:
  """Apply ``patterns`` across a graph, bottom up, and return the rewritten root.

  Iterative over an explicit stack, so depth in the graph never becomes depth on the Python
  stack. Each node is visited once with its already-rewritten children memoized by identity, so
  shared subgraphs stay shared. ``fixpoint`` retries the matcher on a node until nothing fires.
  ``revisit`` instead walks into a replacement's subgraph so nested rewrites collapse in one
  pass. ``max_steps`` bounds the total number of replacements. ``rebuild`` defaults to the
  expression adapter ``rebuild_expr``. Program adapters may reuse ``memo`` across roots with
  identical patterns, rebuild, and options. Expression nodes do not support caller memoization
  because their lowering hints depend on traversal provenance.
  """
  if memo is not None and isinstance(root, Expr):
    raise ValueError("caller memoization is only supported for Program nodes")
  matcher = patterns if isinstance(patterns, PatternMatcher) else PatternMatcher(patterns)
  rebuild = rebuild or cast(Callable[[Node, tuple[Node, ...]], Node], rebuild_expr)
  done: dict[int, Node] = {}
  lowerings: dict[int, Lowering] = {}
  forward: dict[int, Node] = {}
  alive: list[Node] = []
  steps = 0
  stack: list[tuple[Node, bool]] = [(root, False)]
  while stack:
    node, ready = stack.pop()
    if id(node) in done:
      continue
    if memo is not None and node in memo:
      done[id(node)] = memo[node]
      continue
    if not ready:
      stack.append((node, True))
      stack.extend((a, False) for a in reversed(node.args) if id(a) not in done)
      continue
    if (target := forward.pop(id(node), None)) is not None:
      if id(target) not in done:
        raise RuntimeError(f"rewrite replacement at {node.op} contains the node it replaces")
      done[id(node)] = done[id(target)]
      if isinstance(node, Expr):
        lowerings[id(node)] = lowerings[id(target)]
      if memo is not None:
        memo[node] = done[id(node)]
      continue
    args = tuple(done[id(a)] for a in node.args)
    cur = node if all(a is b for a, b in zip(args, node.args, strict=True)) else rebuild(node, args)
    lowering = _combined_lowering(node.lowering, (lowerings[id(arg)] for arg in node.args)) if isinstance(node, Expr) else "auto"
    new = matcher.rewrite(cur)
    while new is not None:
      steps += 1
      if steps > max_steps:
        names = sorted({getattr(p.replacement, "__qualname__", repr(p.replacement)) for p in matcher.candidates(cur.op)})
        raise RuntimeError(f"rewrite exceeded {max_steps} steps at {cur.op} with patterns {names}")
      if isinstance(new, Expr):
        new = cast(Node, _apply_lowering(new, lowering, preserve_identity=id(new) in lowerings))
      alive.append(new)
      if revisit:
        forward[id(node)] = new
        stack.extend(((node, True), (new, False)))
        break
      cur = new
      new = matcher.rewrite(cur) if fixpoint else None
    else:
      done[id(node)] = cur
      if isinstance(node, Expr):
        lowerings[id(node)] = lowering
        lowerings[id(cur)] = lowering
      if memo is not None:
        memo[node] = cur
  return done[id(root)]

def simplify(expr: Expr) -> Expr:
  """Apply algebraic identities and constant folding until the graph stops changing.

  Covers the shared arithmetic identities of ``passes/arith.py`` (neutral elements, zero
  annihilation, ``x - x``, negation normalization, small constant powers), folding of all-constant
  subgraphs, identity reshape, transpose and gather, slice of slice, slice of stack, ``A.T @ v`` as
  ``v @ A`` (and ``v @ A.T`` as ``A @ v``), and a matmul with an all-ones vector as sums.
  """
  return rewrite(expr, SIMPLIFY_PATTERNS, fixpoint=False, revisit=True, max_steps=100_000)

def cse(expr: Expr) -> Expr:
  """Merge structurally equal subgraphs so each distinct computation appears once."""
  return cse_many([expr])[0]

def cse_many(outputs: Iterable[Expr]) -> tuple[Expr, ...]:
  """Common-subexpression elimination across several outputs at once.

  Sharing is found *between* the outputs as well as inside each, which is the point of asking
  for several derivatives from one ``Function.factory`` call.
  """
  outputs = tuple(outputs)
  memo: dict[tuple[Any, ...], Expr] = {}
  replacements: dict[int, Expr] = {}
  for node in topo(outputs):
    cur = _replace_args(node, replacements)
    key = _structural_key(cur)
    cur = memo.setdefault(key, cur)
    replacements[node.id] = cur
  return tuple(replacements[out.id] for out in outputs)