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 sumsqr(x: Any) -> Expr:
  """Sum of squares of every entry: ``dot(x, x)``, a scalar."""
  x = as_expr(x)
  return dot(x, x)

def reduce_max(x: Any, axis: int | tuple[int, ...] | None = None, *, keepdims: bool = False) -> Expr:
  """Maximum over axes, propagating NaN and splitting derivatives equally at ties.

  Axes and ``keepdims`` follow ``Expr.sum``. Structural zeros participate in sparse
  reductions. An empty reduction domain raises ``ValueError`` during construction.
  """
  return _extreme(x, "max", axis, keepdims)

def reduce_min(x: Any, axis: int | tuple[int, ...] | None = None, *, keepdims: bool = False) -> Expr:
  """Minimum over axes, with the same shape, NaN and derivative rules as ``reduce_max``."""
  return _extreme(x, "min", axis, keepdims)

def argmax(x: Any, axis: int | None = None, *, keepdims: bool = False) -> Expr:
  """Index of the first maximum or NaN, as nondifferentiable int64 values.

  Floating operands, a single integer axis or ``None``, and ``keepdims`` follow
  NumPy. With ``axis=None`` indices address the flattened logical array. Structural
  zeros participate. An empty reduction domain raises ``ValueError``.
  Function output leaves currently require float64. Cast indices to float64 to
  return them numerically, or use the integer expression inside the function.
  """
  return _arg_extreme(x, "max", axis, keepdims)

def argmin(x: Any, axis: int | None = None, *, keepdims: bool = False) -> Expr:
  """Index of the first minimum or NaN, with the same rules as ``argmax``."""
  return _arg_extreme(x, "min", axis, keepdims)

def norm_inf(x: Any) -> Expr:
  """Maximum absolute entry, propagating NaN. Empty inputs raise ``ValueError``."""
  return reduce_max(as_expr(x).abs())

def norm_1(x: Any) -> Expr:
  """Sum of absolute entries, with slope zero at zero and value zero for empty inputs."""
  return as_expr(x).abs().sum()

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, *_operands(y, x))

def minimum(x: Any, y: Any) -> Expr:
  """Elementwise minimum, propagating NaN. Derivatives split equally at ties."""
  return binary(ExprOp.MINIMUM, *_operands(x, y))

def maximum(x: Any, y: Any) -> Expr:
  """Elementwise maximum, propagating NaN. Derivatives split equally at ties."""
  return binary(ExprOp.MAXIMUM, *_operands(x, y))

def where(condition: Any, x: Any, y: Any) -> Expr:
  """Choose elementwise: ``x`` where the bool ``condition`` holds, ``y`` elsewhere.

  The three broadcast together and ``x`` and ``y`` share a dtype. The generated C computes both
  branches, so a branch may be NaN or infinite where it is not chosen, and the derivative ignores
  the branch not chosen too: ``where(x > 0, sqrt(x), 0)`` has derivative zero at ``x = 0``,
  including when a ``Function`` or ``vmap`` computes the square root. The condition has no derivative.
  Custom derivative rules receive numerical seeds without condition flags and must return finite
  results on inactive entries. In reverse mode, a select inside a mapped ``Function`` also guards
  derivatives of the map's arguments. A selected zero tangent guards later partials, including across
  ordinary and mapped calls:
  ``sqrt(where(x > 0, x, 0))`` has derivative zero at ``x = -1`` in both modes.
  Dense matrix products drop the condition in reverse mode: for a 2-by-2 matrix of ones ``M``,
  ``where(x > 0, sqrt(x) @ M, 0).sum()`` at ``x = [0, 0]`` gives a zero forward product
  with seed ``[1, 1]``, but gradient ``[NaN, NaN]``.
  """
  condition = _condition(condition)
  x, y = _operands(x, y)
  dtype = elementwise_dtype(ExprOp.SELECT, condition.type.dtype, x.type.dtype, y.type.dtype)
  if condition.type.pattern is not None and x.type.pattern is None and y.type.pattern is None:
    condition = to_dense(condition)
  shape = broadcast_shape(condition.shape, broadcast_shape(x.shape, y.shape))
  type_ = TensorType(
    shape, dtype=dtype, diff=op_diff(ExprOp.SELECT, x, y), pattern=stored_pattern(ExprOp.SELECT, (condition.type, x.type, y.type), shape)
  )
  if type_.pattern is not None:
    chosen = _constant_condition(condition, stored_indices(type_), shape)
    if chosen is not None:
      condition = pack(Expr.const(chosen.reshape(type_.storage_shape), dtype=dtypes.bool_, lowering=condition.lowering), type_.pattern, shape=shape)
  return Expr(
    ExprOp.SELECT,
    (condition, x, y),
    type_,
    lowering=common_lowering(condition, x, y),
  )

def cast(x: Any, dtype: DType | str) -> Expr:
  """Convert ``x`` elementwise to ``dtype``, as NumPy's ``astype`` does.

  A float becomes an integer by dropping its fraction, and a float outside the integer's range,
  NaN included, has no defined result. A cast to ``bool`` is ``x != 0``, so NaN is true. A cast
  keeps its derivative only from one float dtype to another.

  ``uint8`` is not supported for expressions and raises ``ValueError``.
  """
  x, dtype = _condition(x) if isinstance(x, bool) else as_expr(x), as_dtype(dtype)
  if dtype == x.type.dtype:
    return x
  if dtype.is_bool:
    return not_equal(x, 0)
  diff = x.type.diff and x.type.dtype.is_floating and dtype.is_floating
  return Expr(ExprOp.CAST, (x,), TensorType(x.shape, dtype=dtype, diff=diff), lowering=x.lowering)

def equal(x: Any, y: Any) -> Expr:
  """Elementwise ``x == y`` as a bool expression. The ``==`` operator compares expressions by identity."""
  return binary(ExprOp.EQ, *_operands(x, y))

def not_equal(x: Any, y: Any) -> Expr:
  """Elementwise ``x != y`` as a bool expression, true wherever either side is NaN."""
  return binary(ExprOp.NE, *_operands(x, y))

def isfinite(x: Any) -> Expr:
  """Elementwise test that a floating value is neither infinite nor NaN.

  The generated C refuses to compile under ``-ffast-math``, which assumes no value is either.
  """
  return unary(ExprOp.ISFINITE, as_expr(x))

def pack(value: Any, pattern: CanonicalPattern | SparsityPattern, *, shape: tuple[int, ...] | None = None) -> Expr:
  """Build a sparse tensor from its stored values, including explicit zeros.

  A ``SparsityPattern`` takes the values in its own coordinate order, a canonical pattern in
  row-major order. Leading axes of ``value`` are batch axes. ``shape`` defaults to those axes
  followed by the pattern's shape.
  """
  x = as_expr(value)
  if isinstance(pattern, SparsityPattern):
    pattern, permutation = pattern.canonicalize()
    if not x.shape or x.shape[-1] != len(permutation):
      raise ValueError("PACK values' last axis must match the pattern's stored entries")
    indices = np.arange(x.size, dtype=np.int64).reshape(x.shape)[..., permutation]
    x = gather(x, indices)
  if not isinstance(pattern, CanonicalPattern):
    raise TypeError("PACK requires a canonical pattern or SparsityPattern")
  shape = (*x.shape[:-1], *pattern.shape) if shape is None else shape
  type_ = TensorType(shape, x.type.dtype, x.type.diff, pattern)
  return Expr(ExprOp.PACK, (x,), type_, attrs={"pattern": pattern}, lowering=x.lowering)

def values(value: Any) -> Expr:
  """Return a sparse tensor's stored values, in row-major coordinate order, as a dense tensor.

  The result has the storage shape. A dense tensor is returned unchanged.
  """
  x = as_expr(value)
  if x.type.pattern is None:
    return x
  if x.op == ExprOp.PACK:
    return x.args[0]
  return Expr(ExprOp.VALUES, (x,), TensorType(x.type.storage_shape, x.type.dtype, x.type.diff), lowering=x.lowering)

def to_dense(value: Any) -> Expr:
  """Return a dense tensor with zeros outside the pattern. A dense tensor is returned unchanged."""
  x = as_expr(value)
  if x.type.pattern is None:
    return x
  return _scatter_zero(values(x).vec(), stored_indices(x.type), x.size).reshape(x.shape)

def from_dense(value: Any, pattern: CanonicalPattern | SparsityPattern) -> Expr:
  """Keep the entries of a dense tensor inside ``pattern`` as a sparse tensor and drop the others."""
  x = as_expr(value)
  if x.type.pattern is not None:
    raise ValueError("from_dense requires a dense tensor")
  canonical = pattern.canonicalize()[0] if isinstance(pattern, SparsityPattern) else pattern
  type_ = TensorType(x.shape, x.type.dtype, x.type.diff, canonical)
  return pack(gather(x, stored_indices(type_).reshape(type_.storage_shape)), canonical, shape=x.shape)

def sparsify(value: Any) -> Expr:
  """Store only the entries that ``Expr.support`` says may be nonzero.

  Batch axes share one pattern, covering the entries of every batch entry. A sparse tensor is
  returned unchanged.
  """
  x = as_expr(value)
  if x.type.pattern is not None:
    return x
  if len(x.shape) < 2:
    raise ValueError("sparsify requires at least two axes")
  matrix_size = x.shape[-2] * x.shape[-1]
  entries = range(matrix_size) if x.support is None else sorted({i % matrix_size for i in x.support})
  return from_dense(
    x, CanonicalPattern((x.shape[-2], x.shape[-1]), tuple(i // x.shape[-1] for i in entries), tuple(i % x.shape[-1] for i in entries))
  )

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}")