slope.core

   1from pathlib import Path
   2import os
   3import json
   4from typing import (
   5    Callable,
   6    NamedTuple,
   7    Dict,
   8    Hashable,
   9    List,
  10    Any,
  11    Iterable,
  12    Sequence,
  13    Iterator,
  14    Type,
  15    Tuple,
  16    Optional,
  17    Union,
  18    Dict,
  19    Set,
  20    DefaultDict,
  21    Final,
  22)
  23import weakref
  24import types
  25from contextlib import contextmanager, ContextDecorator
  26import itertools
  27import weakref
  28import operator as operator_py
  29import numpy as np
  30import math
  31import inspect
  32from functools import partial, lru_cache
  33import mmap
  34import traceback
  35import importlib
  36import time
  37import cProfile
  38import pstats
  39
  40# =================================
  41#   Utils
  42# =================================
  43
  44
  45class Timing(ContextDecorator):
  46    def __init__(self, prefix="", on_exit=None, enabled=True):
  47        self.prefix, self.on_exit, self.enabled = prefix, on_exit, enabled
  48
  49    def __enter__(self):
  50        self.st = time.perf_counter_ns()
  51
  52    def __exit__(self, *exc):
  53        self.et = time.perf_counter_ns() - self.st
  54        if self.enabled:
  55            print(f"{self.prefix}{self.et*1e-6:6.2f} ms" + (self.on_exit(self.et) if self.on_exit else ""))
  56
  57
  58def colored(st, color: Optional[str], background=False):
  59    return (
  60        f"\u001b[{10*background+60*(color.upper() == color)+30+['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'].index(color.lower())}m{st}\u001b[0m"
  61        if color is not None
  62        else st
  63    )  # replace the termcolor library with one line  # noqa: E501
  64
  65
  66def _format_fcn(fcn):
  67    return f"{fcn[0]}:{fcn[1]}:{fcn[2]}"
  68
  69
  70class Profiling(ContextDecorator):
  71    def __init__(self, enabled=True, sort="cumtime", frac=0.2, fn=None, ts=1):
  72        self.enabled, self.sort, self.frac, self.fn, self.time_scale = enabled, sort, frac, fn, 1e3 / ts
  73
  74    def __enter__(self):
  75        self.pr = cProfile.Profile()
  76        if self.enabled:
  77            self.pr.enable()
  78
  79    def __exit__(self, *exc):
  80        if self.enabled:
  81            self.pr.disable()
  82            if self.fn:
  83                self.pr.dump_stats(self.fn)
  84            stats = pstats.Stats(self.pr).strip_dirs().sort_stats(self.sort)
  85            for fcn in stats.fcn_list[0 : int(len(stats.fcn_list) * self.frac)]:  # type: ignore[attr-defined]
  86                (_primitive_calls, num_calls, tottime, cumtime, callers) = stats.stats[fcn]  # type: ignore[attr-defined]
  87                scallers = sorted(callers.items(), key=lambda x: -x[1][2])
  88                print(
  89                    f"n:{num_calls:8d}  tm:{tottime*self.time_scale:7.2f}ms  tot:{cumtime*self.time_scale:7.2f}ms",
  90                    colored(_format_fcn(fcn), "yellow") + " " * (50 - len(_format_fcn(fcn))),
  91                    colored(f"<- {(scallers[0][1][2]/tottime)*100:3.0f}% {_format_fcn(scallers[0][0])}", "BLACK") if len(scallers) else "",
  92                )
  93
  94
  95def dblog(*msg, enable=True):
  96    if enable:
  97        print(*msg)
  98
  99
 100def unzip2(pairs) -> Tuple[List[Any], List[Any]]:
 101    lst1, lst2 = [], []
 102    for i1, i2 in pairs:
 103        lst1 += [i1]
 104        lst2 += [i2]
 105    return lst1, lst2
 106
 107
 108def list_map(f: Callable, *xs: Iterable) -> List[Any]:
 109    return list(map(f, *xs))
 110
 111
 112def list_zip(*args: List[Any]) -> List[Any]:
 113    fst, *rest = args = list_map(list, args)
 114    n = len(fst)
 115    for arg in rest:
 116        assert len(arg) == n
 117    return list(zip(*args))
 118
 119
 120def split_half(lst: List[Any]) -> Tuple[List[Any], List[Any]]:
 121    assert not len(lst) % 2
 122    return split_list(lst, len(lst) // 2)
 123
 124
 125def merge_lists(which: List[bool], l1: List[Any], l2: List[Any]) -> List[Any]:
 126    l1, l2 = iter(l1), iter(l2)
 127    out = [next(l2) if b else next(l1) for b in which]
 128    assert next(l1, None) is next(l2, None) is None
 129    return out
 130
 131
 132def split_list(lst: List[Any], n: int) -> Tuple[List[Any], List[Any]]:
 133    assert 0 <= n <= len(lst)
 134    return lst[:n], lst[n:]
 135
 136
 137def split_list(lst: List[Any], n: int) -> Tuple[List[Any], List[Any]]:
 138    assert 0 <= n <= len(lst)
 139    return lst[:n], lst[n:]
 140
 141
 142def partition_list(bs: List[bool], l: List[Any]) -> Tuple[List[Any], List[Any]]:
 143    assert len(bs) == len(l)
 144    lists = lst1, lst2 = [], []
 145    for b, x in zip(bs, l):
 146        lists[b].append(x)
 147    return lst1, lst2
 148
 149
 150def lru_cache_verbose(
 151    maxsize: int = 100,
 152    typed: bool = False,
 153    tb_start: int = -12,
 154    tb_end: int = -7,
 155):
 156    def decorator(fn: Callable):
 157        @lru_cache(maxsize=maxsize, typed=typed)
 158        def wrapper(*args, **kwargs) -> Callable:
 159            return fn(*args, **kwargs)
 160
 161        def decorated_function(*args, **kwargs) -> Any:
 162            result = wrapper(*args, **kwargs)
 163            cache_info = wrapper.cache_info()
 164
 165            dblog(
 166                f"{fn.__name__}.{cache_info} {args.__hash__()}",
 167                enable=backend.LOG_LRU,
 168            )
 169            tb = "".join(traceback.format_list(traceback.extract_stack())[tb_start:tb_end]).replace("\n    ", ":\t") + "-" * 20 + "\n"
 170            dblog(f"{tb}", enable=backend.LOG_LRU)
 171
 172            return result
 173
 174        decorated_function.cache_info = wrapper.cache_info
 175        decorated_function.fn = fn
 176        return decorated_function
 177
 178    return decorator
 179
 180
 181def cuda_is_available():
 182    try:
 183        import subprocess
 184        import platform
 185
 186        cmd = f"nvidia-smi{'.exe' if platform.system == 'Windows' else ''}"
 187        result = subprocess.run([cmd], stdout=subprocess.PIPE)
 188        output = result.stdout.decode("utf-8")
 189        return True if "NVIDIA-SMI" in output else False
 190    except FileNotFoundError:
 191        return False
 192
 193
 194class Hashed:
 195    val: Any
 196
 197    def __init__(self, val):
 198        self.val = val
 199
 200    def __hash__(self) -> int:
 201        return hash((self.val,))
 202
 203    def __eq__(self, other):
 204        if isinstance(other, Hashed):
 205            if isinstance(self.val, Tensor) and isinstance(other.val, Tensor):
 206                # because Tensor.__eq__ already for Tensor.equal
 207                return id(self.val) == id(other.val)
 208            return self.val == other.val
 209        return False
 210
 211    def __repr__(self):
 212        return f"Hashed: {repr(self.val)}"
 213
 214
 215# =================================
 216#   Tensors
 217# =================================
 218
 219
 220class DType(NamedTuple):
 221    priority: int
 222    itemsize: int
 223    name: str
 224    mlir: str
 225    numpy: type
 226
 227    @property
 228    def format_code(self):
 229        return f"slope.{self.name}"
 230
 231    def __repr__(self):
 232        return f"<DType: {self.name}>"
 233
 234
 235class dtypes:
 236    float32: Final[DType] = DType(4, 4, "float32", "f32", np.float32)
 237    uint8: Final[DType] = DType(0, 1, "uint8", "u8", np.uint8)
 238    int8: Final[DType] = DType(0, 1, "int8", "i8", np.int8)
 239    bool: Final[DType] = DType(0, 1, "bool", "i1", bool)
 240    int32: Final[DType] = DType(1, 4, "int32", "i32", np.int32)
 241    int64: Final[DType] = DType(2, 8, "int64", "i64", np.int64)
 242    uint64: Final[DType] = DType(2, 8, "uint64", "ui64", np.uint64)
 243    float16: Final[DType] = DType(0, 2, "float16", "f16", np.float16)
 244    half = float16
 245    # bfloat16: Final[DType] = DType(0, 2, "bfloat16", "bf16", np.float16)
 246
 247    all_dtypes = (bool, float16, float32, int8, int32, int64, uint8, uint64)
 248    name_dtype_map = {k.name: k for k in all_dtypes}
 249    name_dtype_map_inv = {v: k for k, v in name_dtype_map.items()}
 250    mlir_dtype_map = {k.mlir: k for k in all_dtypes}
 251    mlir_dtype_map_inv = {v: k for k, v in mlir_dtype_map.items()}
 252
 253    @classmethod
 254    def is_int(cls, dtype):
 255        return dtype in (cls.uint8, cls.int8, cls.int32, cls.uint64, cls.int64)
 256
 257    @classmethod
 258    def is_float(cls, dtype):
 259        return dtype in (cls.float16, cls.bfloat16, cls.float32)
 260
 261
 262class Device(NamedTuple):
 263    name: str
 264    idx: int
 265
 266    @property
 267    def format_code(self):
 268        return f"'{self.name}:{self.idx}'"
 269
 270    def __repr__(self):
 271        return f"<Device: {self.format_code}>"
 272
 273
 274class devices:
 275    cpu: Final[Device] = Device("cpu", 0)
 276    metal: Final[Device] = Device("metal", 0)
 277    cuda0: Final[Device] = Device("cuda", 0)
 278    # TODO: programmatically define this class attrs to support other setup
 279    cuda = cuda0
 280    all_devices = (cpu, metal, cuda0)
 281    name_idx_device_map = {f"{k.name}:{k.idx}": k for k in all_devices}
 282    name_idx_device_map_inv = {v: k for k, v in name_idx_device_map.items()}
 283
 284
 285class TensorBuffer:
 286    def __init__(self, val):
 287        self.val = val
 288
 289
 290class Tensor:
 291    def __init__(self, val: TensorBuffer):
 292        assert isinstance(val, TensorBuffer)
 293        self.buf = val
 294
 295    @property
 296    def symval(self):
 297        return SymbolicTensor.like(self)
 298
 299    @property
 300    def default_dtype(self):
 301        return backend.default_dtype
 302
 303    def is_int(self) -> bool:
 304        return self.dtype in (
 305            dtypes.int8,
 306            dtypes.uint8,
 307            dtypes.uint64,
 308            dtypes.int32,
 309            dtypes.int64,
 310        )
 311
 312    def is_float(self) -> bool:
 313        return self.dtype in (dtypes.float16, dtypes.float32)
 314
 315    def is_unsigned(self) -> bool:
 316        return self.dtype is dtypes.uint8
 317
 318    def to_bool(self):
 319        return self.cast(dtypes.bool)
 320
 321    def short(self):
 322        return self.cast(dtypes.int8)
 323
 324    def int(self):
 325        return self.cast(dtypes.int32)
 326
 327    def long(self):
 328        return self.cast(dtypes.int64)
 329
 330    def half(self):
 331        return self.cast(dtypes.float16)
 332
 333    def float(self):
 334        return self.cast(dtypes.float32)
 335
 336    def __getattr__(self, attr):
 337        if attr in vars(backend.operator_set).keys():
 338            op = getattr(backend.operator_set, attr)
 339            return partial(op, self)
 340        elif attr in vars(backend.procedure_set).keys():
 341            procedure = getattr(backend.procedure_set, attr)
 342            assert not isinstance(procedure, classmethod), f"use {attr} instead of self.{attr}"
 343            return partial(procedure, self)
 344        else:
 345            return self.__getattribute__(attr)
 346
 347    def __getitem__(self, idx):
 348        return self.getitem(idx)
 349
 350    def __setitem__(self, idx, item):
 351        raise NotImplementedError
 352
 353    def str_short(self):
 354        return f"<Tensor: shape={self.shape}, dtype={self.dtype}>"
 355
 356    __neg__ = lambda self: self.neg()
 357    __add__ = lambda self, other: self.add(other)
 358    __radd__ = lambda self, other: self.add(other)
 359    __sub__ = lambda self, other: self.sub(other)
 360    __rsub__ = lambda self, other: self.sub.func(other, self)
 361    __mul__ = lambda self, other: self.mul(other)
 362    __rmul__ = lambda self, other: self.mul(other)
 363    __div__ = lambda self, other: self.div(other)
 364    __rdiv__ = lambda self, other: self.div.func(other, self)
 365    __truediv__ = __div__
 366    __truerdiv__ = __rdiv__
 367    __pow__ = lambda self, other: self.pow(other)
 368    __rpow__ = lambda self, other: self.pow.func(other, self)
 369    __matmul__ = lambda self, other: self.matmul(other)
 370    __rmatmul__ = lambda self, other: self.matmul.func(other, self)
 371    __invert__ = lambda self: self.invert()
 372    __eq__ = lambda self, other: self.equal(other)
 373    __ne__ = lambda self, other: self.not_equal(other)
 374    __ge__ = lambda self, other: self.greater_equal(other)
 375    __le__ = lambda self, other: self.less_equal(other)
 376    __gt__ = lambda self, other: self.greater(other)
 377    __lt__ = lambda self, other: self.less(other)
 378
 379    def __hash__(self):
 380        return id(self.val)
 381
 382    val = property(lambda self: self.buf.val)
 383
 384    def size(self, i):
 385        return self.shape[i]
 386
 387    @property
 388    def dtype(self):
 389        return backend.dtype_of(self)
 390
 391    @property
 392    def device(self):
 393        return backend.device_of(self)
 394
 395    def numpy(self, memmap=False):
 396        return backend.numpy_of(self, memmap)
 397
 398    @property
 399    def shape(self):
 400        return backend.shape_of(self)
 401
 402    @property
 403    def ndim(self):
 404        return len(self.shape)
 405
 406    def numel(self):
 407        return math.prod(self.shape)
 408
 409    def element_size(self):
 410        return self.dtype.itemsize
 411
 412    def nbytes(self):
 413        return self.numel() * self.element_size()
 414
 415    def __repr__(self):
 416        return f"<Tensor: val=\n{self.numpy()}\nshape={self.shape}, dtype={self.dtype.name}, device={self.device.format_code}>"
 417
 418
 419class SymbolicTensor(Tensor):
 420    def __init__(self, shape, dtype, device):
 421        assert isinstance(dtype, DType)
 422        self._shape = tuple(int(i) for i in shape)
 423        self._dtype = dtype
 424        self._device = device
 425
 426    @property
 427    def symval(self):
 428        return self
 429
 430    @property
 431    def val(self):
 432        raise RuntimeError(f"SymbolicTensor actually has no val, from {trace_stack[-1]=}, ")
 433
 434    @property
 435    def shape(self):
 436        return self._shape
 437
 438    @property
 439    def dtype(self):
 440        return self._dtype
 441
 442    @property
 443    def device(self):
 444        return self._device
 445
 446    def like(self, **overrides):
 447        shape = overrides.get("shape", self.shape)
 448        dtype = overrides.get("dtype", self.dtype)
 449        device = overrides.get("device", self.device)
 450        return SymbolicTensor(shape, dtype, device)
 451
 452    def str_short(self):
 453        return f'{str(self.dtype)}[{",".join(str(d) for d in self.shape)}]'
 454
 455    def __hash__(self):
 456        return hash((self.shape, self.dtype))
 457
 458    def __eq__(self, other):
 459        if type(self) != type(other):
 460            return False
 461        return (self.shape == other.shape) and (self.dtype == other.dtype)
 462
 463    def __repr__(self):
 464        return f"<SymbolicTensor: shape={self.shape}, dtype={self.dtype.name}, device={self.device}>"
 465
 466
 467# =================================
 468#   Operator
 469# =================================
 470
 471
 472class Operator:
 473    def __init__(self, name, variadic_inputs=False, nary_outputs=False):
 474        self.name = name
 475        self.variadic_inputs = variadic_inputs
 476        self.nary_outputs = nary_outputs
 477        if self.variadic_inputs:
 478            self.reorg_args = self.reorg_args_nary
 479
 480    def __hash__(self):
 481        return hash(self.name)
 482
 483    def __eq__(self, other):
 484        if not isinstance(other, Operator):
 485            return False
 486        return self.name == other.name
 487
 488    def args_fixer(self, *args, **params):
 489        return args, params
 490
 491    def __call__(self, *args, **params):
 492        args, params = self.reorg_args(args, params)
 493        args, params = self.args_fixer(*args, **params)
 494        ret = bind(self, *args, **params)
 495        if not self.nary_outputs:
 496            ret = ret[0]
 497        return ret
 498
 499    def __repr__(self) -> str:
 500        return f"<{self.name}>"
 501
 502    def typecheck(self, *args, **params):
 503        raise NotImplementedError
 504
 505    def jvp(self, *args, **params):
 506        raise NotImplementedError
 507
 508    def T(self, *args, **params):
 509        raise NotImplementedError
 510
 511    def vmap(self, *args, **params):
 512        raise NotImplementedError
 513
 514    def reorg_args(self, args, params):
 515        sig = inspect.signature(self.typecheck)
 516        args_strs = [k for k, v in sig.parameters.items() if v.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD and k != "self"]
 517        params_strs = [k for k, v in sig.parameters.items() if v.kind == inspect.Parameter.KEYWORD_ONLY and k != "self"]
 518
 519        if args:
 520            if len(args) > len(args_strs):
 521                args, rest = args[: len(args_strs)], args[len(args_strs) :]
 522                if params_strs:
 523                    new_params = {k: rest_arg for k, rest_arg in zip(params_strs, rest) if k not in params}
 524                    params = {**new_params, **params}
 525            else:
 526                args = tuple([params[k] if k in params else arg for k, arg in zip(args_strs, args)])
 527                assert len(args) == len(args_strs)
 528        return args, params
 529
 530    def reorg_args_nary(self, args, params):
 531        return args, params
 532
 533    def partial_run(self, trace, tracers, **params):
 534        tracers_in = [trace.instantiate_const(t) for t in tracers]
 535        symvals_in = [t.symval for t in tracers_in]
 536        symvals_out = self.typecheck(*symvals_in, **params)
 537        tracers_out = [PartialRunTraceTensor(trace, make_unknown_pval(symval), None) for symval in symvals_out]
 538        instruction = InstructionDraft(
 539            self,
 540            tracers_in,
 541            params,
 542            symvals_out,
 543            list_map(weakref.ref, tracers_out),
 544        )
 545        for t in tracers_out:
 546            t.draft = instruction
 547        return tracers_out
 548
 549    def partial_run_instruction(self, unks_in, instruction):
 550        if any(unks_in):
 551            instruction1 = None
 552            instruction2 = Instruction(
 553                instruction.op,
 554                instruction.inputs,
 555                instruction.params,
 556                instruction.out_binders,
 557            )
 558            unks_out = [True for i in instruction.out_binders]
 559            res = [v for unk, v in zip(unks_in, instruction.inputs) if ((not unk) and type(v) is Var)]
 560        else:
 561            instruction1 = instruction
 562            instruction2 = None
 563            unks_out = [False for i in instruction.out_binders]
 564            res = None
 565
 566        return instruction1, instruction2, unks_out, res
 567
 568
 569class MetaOperator(Operator):
 570    def meta_impl(self, *args, **kwargs):
 571        raise NotImplementedError
 572
 573
 574class UnaryOperator(Operator):
 575    def vmap(self, x, *, dim_size, vals_in, dims_in, **params):
 576        (x,), (x_bdim,) = vals_in, dims_in
 577        return [self(x, **params)], [x_bdim]
 578
 579    def typecheck(self, x, **params):
 580        return [SymbolicTensor.like(x)]
 581
 582    def jvp(self, primals, tangents, **params):
 583        (x,), (x_dot,) = primals, tangents
 584        return [self(x, **params)], [self(x_dot, **params)]
 585
 586
 587class BinaryOperator(Operator):
 588    boolean_output = False
 589
 590    def args_fixer(self, x, w, **params):
 591        if isinstance(x, UndefinedPrimal) or type(w) is UndefinedPrimal:
 592            assert x.shape == w.shape
 593            return (x, w), params
 594
 595        if type(x) in TraceTensor.PYTHON_TYPES:
 596            x = backend.full(shape=(), fill_value=x, dtype=w.dtype)
 597        elif type(w) in TraceTensor.PYTHON_TYPES:
 598            w = backend.full(shape=(), fill_value=w, dtype=x.dtype)
 599
 600        shape_delta = x.ndim - w.ndim
 601        if shape_delta > 0:
 602            w = w.reshape((1,) * shape_delta + w.shape)
 603        elif shape_delta < 0:
 604            x = x.reshape((1,) * -shape_delta + x.shape)
 605
 606        shape_ret = tuple([max(x, w) for x, w in zip(x.shape, w.shape)])
 607        if x.shape != shape_ret:
 608            x = x.expand(shape_ret)
 609        if w.shape != shape_ret:
 610            w = w.expand(shape_ret)
 611
 612        if type(x) is Tensor and isinstance(w, TraceTensor):
 613            x = w._trace.pure(x)
 614        elif type(w) is Tensor and isinstance(x, TraceTensor):
 615            w = x._trace.pure(w)
 616        # TODO: https://jax.readthedocs.io/en/latest/type_promotion.html
 617        if x.dtype != w.dtype:
 618            # {int, bool} -> float
 619            if dtypes.is_float(x.dtype) ^ dtypes.is_float(w.dtype):
 620                if dtypes.is_float(w.dtype):
 621                    x = x.cast(w.dtype)
 622                elif dtypes.is_float(x.dtype):
 623                    w = w.cast(x.dtype)
 624            # bool -> int
 625            elif dtypes.is_int(x.dtype) ^ dtypes.is_int(w.dtype):
 626                if dtypes.is_int(w.dtype):
 627                    x = x.cast(w.dtype)
 628                elif dtypes.is_int(x.dtype):
 629                    w = w.cast(x.dtype)
 630            else:  # TODO: fine-grained type promotions
 631                raise NotImplementedError("No other type promotion rules")
 632
 633        return (x, w), params
 634
 635    def vmap(self, dim_size, vals_in, dims_in, **params):
 636        (x, w), (x_bdim, w_bdim) = vals_in, dims_in
 637        if x_bdim != w_bdim:
 638            if x_bdim is None:
 639                x = VMapTrace.move_vmap_dim(x, dim_size, x_bdim, w_bdim)
 640                x_bdim = w_bdim
 641            else:
 642                w = VMapTrace.move_vmap_dim(w, dim_size, w_bdim, x_bdim)
 643        return [self(x, w, **params)], [x_bdim]
 644
 645    def typecheck(self, x: SymbolicTensor, y: SymbolicTensor, **params) -> List[SymbolicTensor]:
 646        if not isinstance(x, (Tensor, SymbolicTensor)) or not isinstance(y, (Tensor, SymbolicTensor)):
 647            raise TypeError
 648        symx = SymbolicTensor.like(x, dtype=dtypes.bool if self.boolean_output else x.dtype)
 649        symy = SymbolicTensor.like(y, dtype=dtypes.bool if self.boolean_output else y.dtype)
 650        if x.dtype != y.dtype:
 651            raise TypeError(f"x.dtype ({x.dtype}) != y.dtype ({y.dtype})")
 652        if symx == symy:
 653            return [symx]
 654        shape_delta = len(symx.shape) - len(symy.shape)
 655        if shape_delta > 0:
 656            symy = symy.like(shape=(1,) * shape_delta + symy.shape)
 657        elif shape_delta < 0:
 658            symx = symx.like(shape=(1,) * -shape_delta + symx.shape)
 659        if symx == symy:
 660            return [symx]
 661        else:
 662            shape_ret = tuple([max(x, w) for x, w in zip(symx.shape, symy.shape)])
 663            if symx.shape != shape_ret:
 664                symx = symx.like(shape=shape_ret)
 665            if symy.shape != shape_ret:
 666                symy = symx.like(shape=shape_ret)
 667            if symx != symy:
 668                raise TypeError(f"symx ({symx}) != symy ({symy})")
 669            return [symx]
 670
 671    def jvp(self, primals, tangents, **params):
 672        (x, w), (x_dot, w_dot) = primals, tangents
 673        return [self(x, w, **params)], [self(x_dot, w_dot, **params)]
 674
 675    def T(self, cotangents, x, w):
 676        (gL_y,) = cotangents
 677        if self.boolean_output:
 678            gL_y = gL_y.cast(x.dtype)
 679        if isinstance(x, UndefinedPrimal):
 680            return [gL_y, NullCotangent]
 681        elif isinstance(w, UndefinedPrimal):
 682            return [NullCotangent, gL_y]
 683        else:
 684            raise ValueError
 685
 686
 687class ReduceOperator(Operator):
 688    def args_fixer(self, x, *, dim=None, keepdim=False):
 689        if dim is None:
 690            dim = tuple(range(x.ndim))
 691        elif isinstance(dim, int):
 692            dim = (dim,)
 693        dim = tuple(a if a >= 0 else a + len(x.shape) for a in dim)
 694        return (x,), dict(dim=dim, keepdim=keepdim)
 695
 696    def vmap(self, dim_size, vals_in, dims_in, *, dim, keepdim):
 697        (x,), (x_bdim,) = vals_in, dims_in
 698        dim = tuple(a + (x_bdim <= a) for a in dim)
 699        out_bdim = x_bdim - sum(a < x_bdim for a in dim)
 700        return [self(x, dim=dim, keepdim=keepdim)], [out_bdim]
 701
 702    def typecheck(self, x: SymbolicTensor, *, dim=None, keepdim=False) -> List[SymbolicTensor]:
 703        dim = list(set([a + len(x.shape) if a < 0 else a for a in dim]))
 704        if keepdim:
 705            new_shape = [d if i not in dim else 1 for i, d in enumerate(x.shape)]
 706        else:
 707            new_shape = [d for i, d in enumerate(x.shape) if i not in dim]
 708        return [SymbolicTensor.like(x, shape=tuple(new_shape))]
 709
 710
 711class InitOperator(Operator):
 712    def vmap(self, dim_size, vals_in, dims_in, **params):
 713        (x_bdim,) = dims_in
 714        y = self(**params)
 715        y = y.unsqueeze(x_bdim)
 716        return [y], [x_bdim]
 717
 718    def jvp(self, primals, tangents, **params):
 719        y = self(**params)
 720        y_dot = NullCotangent(y.symval)
 721        return [y], [y_dot]
 722
 723    def T(self, cotangents, **params):
 724        return [NullCotangent(cotangents[0])]
 725
 726
 727class ShapeOperator(Operator):
 728    pass
 729
 730
 731class GeneralReduceOperator(Operator):
 732    pass
 733
 734
 735class OperatorSet:
 736    def __init__(self):
 737        self.register("jit_op")(JitOp)
 738
 739    def register(self, name, variadic_inputs=False, nary_outputs=False, aliases=()):
 740        def wrap(op_cls):
 741            assert name not in vars(self)
 742            op = op_cls(name, variadic_inputs, nary_outputs)
 743            setattr(self, name, op)
 744            for a in aliases:
 745                setattr(self, a, op)
 746            return op_cls
 747
 748        return wrap
 749
 750
 751class ProcedureSet:
 752    def register(self, aliases=()):
 753        def wrap(f):
 754            assert f.__name__ not in vars(self)
 755            setattr(self, f.__name__, f)
 756            for a in aliases:
 757                setattr(self, a, f)
 758            return f
 759
 760        return wrap
 761
 762
 763class CodegenOutput(NamedTuple):
 764    code_lines: List[str]
 765    fn_defs: Dict[str, List[str]]
 766    in_binders: List["ProgramEnvVar"]
 767    outs: List["ProgramEnvVar"]
 768
 769
 770class Backend:
 771    LOG_LRU = int(os.environ.get("LOG_LRU", 0))
 772    LOG_JIT = int(os.environ.get("LOG_JIT", 0))
 773    LOG_TREE = int(os.environ.get("LOG_TREE", 0))
 774    LOG_BACKEND = int(os.environ.get("LOG_BACKEND", 0))
 775    LOG_PROGRAM = int(os.environ.get("LOG_PROGRAM", 0))
 776    LOG_INIT = int(os.environ.get("LOG_INIT", 1))
 777    device_var = os.environ.get("DEFAULT_DEVICE", "cpu:0")
 778    if device_var[-2] != ":":
 779        device_var += ":0"
 780    DEFAULT_DEVICE = devices.name_idx_device_map[device_var]
 781    DEFAULT_DTYPE = dtypes.name_dtype_map[os.environ.get("DEFAULT_DTYPE", "float32")]
 782    dtype_for_indices: DType = None  # need to override
 783
 784    def __init__(
 785        self,
 786        operator_set: OperatorSet,
 787        procedure_set: ProcedureSet,
 788    ):
 789        self.operator_set = operator_set
 790        self.procedure_set = procedure_set
 791        self.node_types = dict()
 792        self.impls = dict()
 793        self.register_node(tuple, lambda t: (None, t), lambda _, xs: tuple(xs), "tuple")
 794        self.register_node(list, lambda l: (None, l), lambda _, xs: list(xs), "list")
 795        self.register_node(
 796            dict,
 797            lambda d: list_map(tuple, unzip2(sorted(d.items()))),
 798            lambda keys, vals: dict(list_zip(keys, vals)),
 799            "dict",
 800        )
 801        self.register_node(
 802            UndefinedPrimal,
 803            lambda u: (u.symval, ()),
 804            lambda symval, _: UndefinedPrimal(symval),
 805            "UndefinedPrimal",
 806        )
 807
 808    def set_impl(self, op: Union[types.LambdaType, types.FunctionType]):
 809        def set_impl_(fn):
 810            self.impls[op] = types.MethodType(fn, self)
 811
 812        return set_impl_
 813
 814    def register_node(self, ty: Type, to_iter: Callable, from_iter: Callable, name=None) -> None:
 815        if name is None:
 816            name = str(ty)
 817        self.node_types[ty] = NodeType(name, to_iter, from_iter)
 818
 819    def __getattr__(self, attr):
 820        try:
 821            dblog(
 822                f"Looking {self}.{attr} in operator_set",
 823                enable=backend.LOG_BACKEND,
 824            )
 825            return getattr(self.operator_set, attr)
 826        except:
 827            pass
 828        try:
 829            dblog(
 830                f"Looking {self}.{attr} in procedure_set",
 831                enable=backend.LOG_BACKEND,
 832            )
 833            return getattr(self.procedure_set, attr)
 834        except:
 835            pass
 836        dblog(
 837            f"Fallback to default {self} getattribute",
 838            enable=backend.LOG_BACKEND,
 839        )
 840        super().__getattribute__(attr)
 841
 842    def tensor(
 843        self,
 844        val: Union[list, tuple, np.ndarray, "TensorBuffer"] = None,
 845        dtype: Optional[Any] = None,
 846        device=None,
 847    ):
 848        if isinstance(val, TensorBuffer):
 849            return Tensor(val)
 850        elif isinstance(val, Tensor):
 851            return val
 852        if type(val) is bytes:
 853            val = np.frombuffer(val, dtype=dtype)
 854        return self.from_numpy(val, dtype, device)
 855
 856    def symbolic_tensor(
 857        self,
 858        shape: Union[list, tuple, np.ndarray, "TensorBuffer"] = None,
 859        dtype: Optional[Any] = None,
 860        device=None,
 861    ):
 862        dtype = dtype or self.DEFAULT_DTYPE
 863        device = device or self.DEFAULT_DEVICE
 864        return SymbolicTensor(shape, dtype, device)
 865
 866    def seed(self, seed):
 867        raise NotImplementedError
 868
 869    @property
 870    def default_dtype_value(self):
 871        return self.dtype_map[backend.DEFAULT_DTYPE]
 872
 873    def set_method(self, method):
 874        setattr(self, method.__name__, types.MethodType(method, self))
 875
 876    def from_numpy(self, val, device):
 877        raise NotImplementedError
 878
 879    def numpy_of(self, tensor):
 880        raise NotImplementedError
 881
 882    def device_of(self, tensor):
 883        raise NotImplementedError
 884
 885    def shape_of(self, tensor):
 886        raise NotImplementedError
 887
 888    def dtype_of(self, tensor):
 889        raise NotImplementedError
 890
 891    @lru_cache_verbose()
 892    def jit_program(
 893        self,
 894        hashed_program: Hashed,
 895        hashed_consts: Tuple[Hashed, ...],
 896    ):
 897        program: Program = hashed_program.val
 898        typecheck_program(program)
 899        consts = [x.val for x in hashed_consts]
 900        in_symvals = [v.symval for v in program.in_binders[len(consts) :]]
 901        codegen_output: CodegenOutput = self.codegen(program, consts + in_symvals, fn_name="main")
 902        fn, code = self.compile(codegen_output)
 903        jit_output = JitOutput(program, codegen_output, fn, code, consts)
 904        return jit_output
 905
 906    def codegen(self, program: "Program", args: Tuple, in_symvals: Tuple, name: str):
 907        "Returns compiler IR from the Program"
 908        raise NotImplementedError
 909
 910    def compile(self, program: "Program", args: Tuple, in_symvals: Tuple, name: str):
 911        "Compiles compiler IR to a Python callable function"
 912        raise NotImplementedError
 913
 914    def export(self, jit_output, *args, **params):
 915        raise NotImplementedError
 916
 917    def load(self, path, single_key="_tensor"):
 918        with open(path, mode="rb") as f:
 919            with mmap.mmap(f.fileno(), length=0, access=mmap.ACCESS_READ) as m:
 920                json_len = np.int64(m[0])
 921                start = 8 + json_len
 922                metadata = json.loads(m[8:start])
 923                ret = {}
 924                for k, v in metadata.items():
 925                    if k != "__metadata__":
 926                        dtype = Tensor.mlir_dtype_map[(v["dtype"])]
 927                        data_start = start + v["data_offsets"][0]
 928                        data_end = start + v["data_offsets"][1]
 929                        t_np = np.frombuffer(m[data_start:data_end], dtype=dtype.numpy())
 930                        t = backend.tensor(t_np, dtype=dtype)
 931                        t = t.reshape(tuple(v["shape"]))
 932                        ret[k] = t
 933                if len(ret) == 1 and single_key in ret.keys():
 934                    return ret[single_key]
 935                return ret
 936
 937    def save(self, tensors: Dict[str, Tensor], path: str, single_key="_tensor"):
 938        if isinstance(tensors, Tensor):
 939            tensors = {single_key: tensors}
 940        else:
 941            assert all((isinstance(k, str) and isinstance(v, Tensor)) for k, v in tensors.items())
 942
 943        metadata, offset = {}, 0
 944        for k, v in tensors.items():
 945            metadata[k] = {
 946                "dtype": v.dtype.mlir,
 947                "shape": list(v.shape),
 948                "data_offsets": [offset, offset + v.nbytes()],
 949            }
 950            offset += v.nbytes()
 951        j = json.dumps(metadata, separators=(",", ":"))
 952        Path(path).unlink(missing_ok=True)
 953        jbytes = j.encode("utf-8")
 954        start = 8 + len(jbytes)
 955        with open(path, mode="wb") as f:  # make empty file, fill with enough space
 956            f.write(b"\x00" * (start + offset))
 957        with open(path, mode="r+b") as f:
 958            with mmap.mmap(f.fileno(), length=0, access=mmap.ACCESS_WRITE) as m:
 959                m[0:8] = np.int64(len(j)).tobytes()
 960                m[8:start] = jbytes
 961                for t, tm in zip(tensors.values(), metadata.values()):
 962                    data_start, data_end = tm["data_offsets"]
 963                    m[start + data_start : start + data_end] = t.numpy().tobytes()
 964
 965
 966# =================================
 967#   Program
 968# =================================
 969
 970
 971class Var:
 972    def __init__(self, symval):
 973        self.symval = symval
 974        self.val = None
 975
 976
 977class Lit:
 978    def __init__(self, val):
 979        self.symval = SymbolicTensor.like(get_symval(val))
 980        self.val = val
 981
 982
 983Atom = Union[Var, Lit]
 984
 985
 986class Instruction(NamedTuple):
 987    op: Operator
 988    inputs: List[Atom]
 989    params: Dict[str, Any]
 990    out_binders: List[Atom]
 991
 992
 993class ProgramEnvVar(NamedTuple):
 994    name: str
 995    symval: SymbolicTensor
 996    is_const: bool = False
 997
 998    @property
 999    def shape(self):
1000        return self.symval.shape
1001
1002    @property
1003    def dtype(self):
1004        return self.symval.dtype
1005
1006    @property
1007    def device(self):
1008        return self.symval.device
1009
1010    @property
1011    def ndim(self):
1012        return self.symval.ndim
1013
1014    def numpy(self):
1015        return self.symval.numpy()
1016
1017    def __repr__(self):
1018        return f"<ProgramEnvVar: name={self.name}, symval={self.symval}>"
1019
1020    str_short = __repr__
1021
1022
1023class Program:
1024    def __init__(
1025        self,
1026        in_binders: Any,
1027        instructions: Tuple[Instruction],
1028        outs: Any,
1029        num_consts: int = 0,
1030        static_args: Any = (),
1031        name: str = "my_program",
1032        indent_amount=4,
1033    ):
1034        self.in_binders: Any = in_binders
1035        self.outs: Any = outs
1036        self.instructions = self.prune_instructions(instructions, outs)
1037        self.num_consts: int = num_consts
1038        self.static_args = static_args
1039        self.name: str = name
1040        self.indent_amount: int = indent_amount
1041
1042        self.env: Dict[ProgramEnvVar, Any] = dict()
1043        for inb in self.in_binders:
1044            prefix = "x" if type(inb.symval) is SymbolicTensor else "c"
1045            idx = sum([1 if v.name[0] == prefix else 0 for v in self.env.values()])
1046            self.env[inb] = ProgramEnvVar(f"{prefix}{idx}", inb.symval, True if prefix == "c" else False)
1047        for instruction in self.instructions:
1048            if len(instruction.out_binders) == 0:
1049                continue
1050            for outb in instruction.out_binders:
1051                prefix = "y" if outb in self.outs else "z"
1052                idx = sum([1 if v.name[0] == prefix else 0 for v in self.env.values()])
1053                self.env[outb] = ProgramEnvVar(f"{prefix}{idx}", outb.symval)
1054        self.curr_repr = repr(self)
1055
1056    def pprint_shape(self, symval, scalar_as_empty_array=False):
1057        xdtype = symval.dtype.mlir
1058        if len(symval.shape) > 0:
1059            xshape = f"{', '.join((repr(i) for i in symval.shape))}"
1060            return f"[{xshape}, {xdtype}]"
1061        else:
1062            return f"[{xdtype}]"
1063
1064    def pprint_sig(self, in_symvals, out_symvals, unpack_unary_output=False):
1065        in_code = ", ".join(self.pprint_shape(t) for t in in_symvals)
1066        in_code = f"({in_code})" if len(in_symvals) > 1 else in_code
1067        out_code = ", ".join(self.pprint_shape(t) for t in out_symvals)
1068        out_code = f"({out_code})" if len(out_symvals) > 1 or unpack_unary_output else out_code
1069        typing_code = f"{in_code} -> {out_code}"
1070        return typing_code
1071
1072    def __repr__(self):
1073        fn_defs = self.instructions_as_code(self, dict())
1074        return "\n".join(line for code_lines in fn_defs.values() for line in code_lines)
1075
1076    def save(self, *args, dir_path="/tmp/slope_program", dry_run=False):
1077        os.makedirs(dir_path, exist_ok=True)
1078        head_code_lines = [f"import slope # backend={backend.__class__.__name__}"]
1079        fn_defs = self.instructions_as_code(self, dict())
1080        in_binders_vars = [self.env[i] for i in self.in_binders]
1081        for i in range(len(self.in_binders)):
1082            ibv = in_binders_vars[i]
1083            if ibv.is_const:
1084                const_filename = f"{ibv.name}.safetensors"
1085                const_path = os.path.join(dir_path, f"{const_filename}")
1086                if not dry_run:
1087                    backend.save(args[i], const_path)
1088                dblog(
1089                    f"Saved {ibv.name} at {const_path}",
1090                    enable=backend.LOG_BACKEND,
1091                )
1092                head_code_lines += [f"""{ibv.name} = slope.load("./{const_filename}")"""]
1093        head_code_lines += [""]
1094        code = "\n".join(head_code_lines + [line for code_lines in fn_defs.values() for line in code_lines])
1095        dblog(
1096            f"Contents of {self.name}:\n```\n{code}\n```",
1097            enable=backend.LOG_BACKEND,
1098        )
1099        program_path = os.path.join(dir_path, "main.py")
1100        if not dry_run:
1101            with open(program_path, "w") as f:
1102                f.write(code)
1103        dblog(
1104            f"Saved program {self.name} at {program_path}",
1105            enable=backend.LOG_BACKEND,
1106        )
1107        ls_contents = "\n\t".join(os.listdir(dir_path))
1108        dblog(
1109            f"Contents of {dir_path}:\n\t{ls_contents}",
1110            enable=backend.LOG_BACKEND,
1111        )
1112
1113    def __hash__(self):
1114        return hash(self.curr_repr)
1115
1116    def __eq__(self, other):
1117        return self is other
1118
1119    @classmethod
1120    def instructions_as_code(cls, program, fn_defs):
1121        def indent(code, indent_amount):
1122            spaces = " " * (len(code) - len(code.lstrip()))
1123            spaces += " " * indent_amount
1124            return "\n".join([spaces + line for line in code.strip().split("\n")])
1125
1126        in_binders_vars = [program.env[i] for i in program.in_binders]
1127        body_code_lines = []
1128        for instruction in program.instructions:
1129            if len(instruction.out_binders) == 0:
1130                continue
1131            params = instruction.params.copy()
1132            for param_name, param in params.items():
1133                if isinstance(param, Program):
1134                    sub_program = param
1135                    fn_defs = cls.instructions_as_code(sub_program, fn_defs)
1136                    program_in_vals = ", ".join(f"{program.env[x].name}" for x in instruction.inputs)
1137                    params[param_name] = f"slope.make_program({sub_program.name}, {program_in_vals})[0]"
1138                if isinstance(param, DType):
1139                    params[param_name] = f"slope.{param.name}"
1140            param_vals = ", ".join(f"{param_name}={param}" for param_name, param in params.items())
1141            in_vals = ", ".join(f"{program.env[x].name}" for x in instruction.inputs)
1142            out_vals = ", ".join(f"{program.env[z].name}" for z in instruction.out_binders)
1143            sig = program.pprint_sig(
1144                [program.env[x].symval for x in instruction.inputs],
1145                [program.env[y].symval for y in instruction.out_binders],
1146            )
1147            line = f"""{out_vals} = slope.{instruction.op.name}({in_vals}{", " if (param_vals and in_vals) else ""}{param_vals}) # {sig}"""
1148            body_code_lines += [indent(line, program.indent_amount)]
1149
1150        fn_args_str = ", ".join([f"{i.name}" for i in in_binders_vars])
1151        # fn_static_args_str = ", ".join([f"{a}={a_val}" for a, a_val in program.static_args])
1152        out_vars = [program.env[o] for o in program.outs]
1153        fn_sig = program.pprint_sig(
1154            [i.symval for i in in_binders_vars],
1155            [o.symval for o in out_vars],
1156        )
1157        head_code_line = [f"def {program.name}({fn_args_str}): # {fn_sig}"]
1158        out_str = ", ".join([f"{o.name}" for o in out_vars])
1159        tail_code_line = [indent(f"return {out_str}", program.indent_amount)]
1160        code_lines = head_code_line + body_code_lines + tail_code_line + ["\n"]
1161
1162        fn_defs[program.name] = code_lines
1163        return fn_defs
1164
1165    @staticmethod
1166    def prune_instructions(instructions, outs):
1167        graph = dict()
1168        for instruction in instructions:
1169            parent_nodes, child_nodes = instruction.out_binders, instruction.inputs
1170            for parent in parent_nodes:
1171                if parent not in graph:
1172                    graph[parent] = set()
1173                for child in child_nodes:
1174                    graph[parent].add(child)
1175        visited_from_terminal = set()
1176
1177        def dfs(node, visited):
1178            visited.add(node)
1179            if node in graph:
1180                for neighbor in graph[node]:
1181                    if neighbor not in visited:
1182                        dfs(neighbor, visited)
1183
1184        for terminal_node in outs:
1185            dfs(terminal_node, visited_from_terminal)
1186        unreachable_nodes = set(graph.keys()) - visited_from_terminal
1187
1188        instructions_to_prune = []
1189        for instruction in instructions:
1190            parent_nodes, child_nodes = instruction.out_binders, instruction.inputs
1191            if any(node in unreachable_nodes for node in parent_nodes) or any(node in unreachable_nodes for node in child_nodes):
1192                instructions_to_prune += [instruction]
1193        new_instructions = [inst for inst in instructions if inst not in instructions_to_prune]
1194        if backend.LOG_PROGRAM:
1195            LI = len(instructions)
1196            LNI = len(new_instructions)
1197            DIFF = LI - LNI
1198            UN = len(unreachable_nodes)
1199            dblog(f"Before: {LI}\tAfter: {LNI}\tDiff vs Unreachables: {DIFF} == {UN} = {DIFF==UN}")
1200        return new_instructions
1201
1202
1203class ProgramType(NamedTuple):
1204    in_types: Tuple[SymbolicTensor]
1205    out_types: Tuple[SymbolicTensor]
1206
1207    def __repr__(self):
1208        in_types = ", ".join(symval.str_short() for symval in self.in_types)
1209        out_types = ", ".join(symval.str_short() for symval in self.out_types)
1210        return f"({in_types}) -> ({out_types})"
1211
1212
1213# =================================
1214#   Tracer and Trace
1215# =================================
1216
1217
1218class Empty:
1219    pass
1220
1221
1222empty = Empty()
1223
1224
1225class Store:
1226    val = empty
1227
1228    def set_value(self, val):
1229        assert self.val is empty
1230        self.val = val
1231
1232    def __call__(self):
1233        return self.val
1234
1235
1236class NodeType(NamedTuple):
1237    name: str
1238    flatten: Callable
1239    unflatten: Callable
1240
1241
1242class TreeDef(NamedTuple):
1243    node_type: NodeType
1244    node_metadata: Hashable
1245    child_treedefs: Tuple["TreeDef", ...]
1246
1247    def __repr__(self):
1248        ret = self.tree_repr(self)
1249        return ret
1250
1251    def tree_repr(self, tree, indent="  ", prefix="", last=True):
1252        ret = ""
1253
1254        def _tree_repr(tree, indent, prefix, last):
1255            nonlocal ret
1256            if isinstance(tree, TreeDef):
1257                ret += f'{prefix} {("└─" if last else "├─")} {tree.node_type.name}\n'
1258                for i, item in enumerate(tree.child_treedefs):
1259                    new_prefix = prefix + (indent if not last else "   ")
1260                    new_last = i == len(tree.child_treedefs) - 1
1261                    _tree_repr(item, indent, new_prefix, new_last)
1262            else:
1263                ret += f'{prefix} {("└─" if last else "├─")} {tree}\n'
1264
1265        _tree_repr(tree, indent="  ", prefix="", last=True)
1266        return ret
1267
1268    @property
1269    def num_leaves(self):
1270        def get_num_leaves(x):
1271            if isinstance(x, Leaf):
1272                return 1
1273            else:
1274                return sum(get_num_leaves(sub_x) for sub_x in x.child_treedefs)
1275
1276        return sum(get_num_leaves(x) for x in self.child_treedefs)
1277
1278
1279class Leaf:
1280    def __init__(self, val):
1281        if hasattr(val, "shape"):
1282            val = SymbolicTensor.like(val)
1283        self.val = val
1284
1285    def __repr__(self):
1286        ret = self.val.str_short() if isinstance(self.val, SymbolicTensor) else repr(self.val)
1287        return f"<Leaf: {ret}>"
1288
1289    def __hash__(self):
1290        return hash(self.val)
1291
1292    def __eq__(self, other):
1293        return True  # make TreeDef __eq__ don't care Leaf
1294        # if isinstance(other, Leaf): # TODO: test above assumption
1295        #     return self.val == other.val
1296
1297
1298# =================================
1299#   jit operator
1300# =================================
1301
1302
1303class JitOutput:
1304    def __init__(self, program: Program, codegen_output: CodegenOutput, fn, code: str, consts: List[Any]):
1305        super().__init__()
1306        self.program = program
1307        self.code = code
1308        self.codegen_output = codegen_output
1309        self.fn: Callable = fn
1310        self.consts = consts
1311
1312    def __call__(self, *args, **params):
1313        args, in_tree = tree_flatten(args)
1314        args = tree_map(lambda a: a.val if isinstance(a, Tensor) else a, args)
1315        try:
1316            outs = self.fn(*args, **params)
1317            if not isinstance(outs, tuple):  # TODO: IREE FunctionInvoker destructure 1-tuple, need to undo
1318                outs = (outs,)
1319        except Exception as e:
1320            dblog(self.code, enable=backend.LOG_JIT)
1321            raise
1322        return [backend.tensor(TensorBuffer(o)) for o in outs]
1323
1324
1325class JitOp(MetaOperator):
1326    def meta_impl(self, *args, program: Program, **_):
1327        hashed_program = Hashed(program)
1328        num_consts = program.num_consts
1329        consts, args = args[:num_consts], args[num_consts:]
1330        hashed_consts = tuple(map(Hashed, consts))
1331        jit_output = backend.jit_program(hashed_program, hashed_consts)
1332        ret = jit_output(*consts, *args)
1333        return ret
1334
1335    def reorg_args(self, args, params):
1336        return args, params
1337
1338    def typecheck(self, *in_types, program: Program):
1339        program_type = typecheck_program(program)
1340        if not all(t1 == t2 for t1, t2 in zip(program_type.in_types, in_types)):
1341            ret = "Type mismatch program.in_types vs in_types:\n"
1342            for i, j in zip(program_type.in_types, in_types):
1343                ret += f"{i}, {j}, {i == j}"
1344            raise TypeError(ret)
1345        return program_type.out_types
1346
1347    def vmap(self, dim_size, vals_in, dims_in, program: Program):
1348        program, consts = vmap_program(program, dim_size, tuple(dims_in))
1349        outs = self(*consts, *vals_in, program=program)
1350        if not isinstance(outs, tuple):
1351            outs = (outs,)
1352        return outs, [0] * len(outs)
1353
1354    def jvp(self, primals, tangents, *, program):
1355        new_program, new_consts = jvp_program(program)
1356        outs = bind(
1357            self,
1358            *new_consts,
1359            *primals,
1360            *tangents,
1361            program=new_program,
1362        )
1363        n = len(outs) // 2
1364        primals_out, tangents_out = outs[:n], outs[n:]
1365        return primals_out, tangents_out
1366
1367    def T(self, cotangents, *invals, program):
1368        undef_primals = [isinstance(x, UndefinedPrimal) for x in invals]
1369        transposed_program, new_consts = transpose_program(program, tuple(undef_primals))
1370
1371        residuals, _ = partition_list(undef_primals, invals)
1372        outs = bind(
1373            self,
1374            *new_consts,
1375            *residuals,
1376            *cotangents,
1377            program=transposed_program,
1378        )
1379        outs = iter(outs)
1380
1381        return [next(outs) if undef else None for undef in undef_primals]
1382
1383    def partial_run(self, trace, tracers, *, program):
1384        in_unknowns = [not t.pval.is_known for t in tracers]
1385        program1, program2, out_unknowns, num_res = partial_run_program(program, in_unknowns)
1386        known_tracers, unknown_tracers = partition_list(in_unknowns, tracers)
1387        known_vals = [t.pval.const for t in known_tracers]
1388        outs1_res = bind(backend.jit_op, *known_vals, program=program1)
1389        outs1, res = split_list(outs1_res, len(program1.outs) - num_res)
1390        res_tracers = [trace.instantiate_const(full_raise(trace, x)) for x in res]
1391        outs2 = [PartialRunTraceTensor(trace, make_unknown_pval(v.symval), None) for v in program2.outs]
1392        instruction = InstructionDraft(
1393            self,
1394            res_tracers + unknown_tracers,
1395            dict(program=program2),
1396            [v.symval for v in program2.outs],
1397            list_map(weakref.ref, outs2),
1398        )
1399        for t in outs2:
1400            t.draft = instruction
1401
1402        return merge_lists(out_unknowns, outs1, outs2)
1403
1404    def partial_run_instruction(self, unks_in, instruction) -> Tuple[Instruction, Instruction, List[bool], List[Var]]:
1405        program = instruction.params["program"]
1406        program1, program2, out_unknowns, num_res = partial_run_program(program, unks_in)
1407        ins1, ins2 = partition_list(unks_in, instruction.inputs)
1408        out_binders1, out_binders2 = partition_list(out_unknowns, instruction.out_binders)
1409        res = [Var(v.symval) for v in program2.in_binders[:num_res]]
1410        instruction1 = Instruction(self, ins1, dict(program=program1), out_binders1 + res)
1411        instruction2 = Instruction(self, res + ins2, dict(program=program2), out_binders2)
1412        return instruction1, instruction2, out_unknowns, res
1413
1414
1415class MainTrace(NamedTuple):
1416    level: int
1417    trace_type: Type["Trace"]
1418    global_data: Optional[Any]
1419
1420
1421class Trace:
1422    main: MainTrace
1423
1424    def __init__(self, main: MainTrace) -> None:
1425        self.main = main
1426
1427    def pure(self, val):
1428        raise NotImplementedError
1429
1430    def run_op(self, op, tracers, params):
1431        raise NotImplementedError
1432
1433
1434class RunTrace(Trace):
1435    pure = lambda self, x: x
1436
1437    def run_op(self, op: Operator, args, params):
1438        if isinstance(op, MetaOperator):
1439            args, params = op.reorg_args(args, params)
1440            args, params = op.args_fixer(*args, **params)
1441            ret = op.meta_impl(*args, **params)
1442        else:
1443            fn = self.get_fn(op, *tuple(SymbolicTensor.like(a) for a in args), **params)
1444            # with Timing(f"RUN {op}"):ret = jit(
1445            ret = jit(
1446                fn,
1447                static_argnames=("params",),
1448                name=jit.get_jit_name(args, params, op.name),
1449            )(*args, **params)
1450
1451        return ret
1452
1453    @staticmethod
1454    @lru_cache_verbose()
1455    def get_fn(op, *symval_args, **params):
1456        def fn(*args, **params):
1457            return [op(*args, **params)]
1458
1459        return fn
1460
1461
1462class SymbolicRunTrace(Trace):
1463    # pure = lambda self, x: x
1464    def pure(self, val: Any) -> SymbolicTensor:
1465        return val.symval
1466
1467    def run_op(self, op, tracers, params):
1468        symvals_in = tree_map(lambda x: x.symval, tracers)
1469        symvals_out = op.typecheck(*symvals_in, **params)
1470        return symvals_out
1471
1472
1473class TraceTensor(Tensor):
1474    PYTHON_TYPES = {
1475        bool,
1476        int,
1477        float,
1478    }
1479    _trace: "Trace"
1480
1481    def __init__(self, *args, **kwargs):
1482        raise NotImplementedError
1483
1484    symval = property(lambda self: get_symval(self.val))
1485    dtype = property(lambda self: self.symval.dtype)
1486    shape = property(lambda self: self.symval.shape)
1487    device = property(lambda self: self.symval.device)
1488
1489    @property
1490    def val(self):
1491        raise NotImplementedError
1492
1493    def __str__(self):
1494        return repr(self)
1495
1496    def full_lower(self):
1497        return self
1498
1499    @property
1500    def ndim(self):
1501        return len(self.shape)
1502
1503    def __repr__(self):
1504        return f"{self.__class__.__name__}({repr(self.symval)})"
1505
1506
1507class VMapTraceTensor(TraceTensor):
1508    def __init__(self, trace, val, vmap_dim):
1509        self._trace = trace
1510        self._val = val
1511        self.vmap_dim = vmap_dim
1512
1513    @property
1514    def val(self):
1515        return self._val
1516
1517    @property
1518    def symval(self):
1519        symval = get_symval(self.val)
1520        if self.vmap_dim is None:
1521            return symval
1522        else:
1523            shape = list(symval.shape)
1524            del shape[self.vmap_dim]
1525            return symval.like(shape=tuple(shape))
1526
1527    def full_lower(self):
1528        if self.vmap_dim is None:
1529            return full_lower(self.val)
1530        else:
1531            return self
1532
1533
1534class VMapTrace(Trace):
1535    pure = lambda self, val: VMapTraceTensor(self, val, None)
1536
1537    @property
1538    def dim_size(self):
1539        return self.main.global_data
1540
1541    def run_op(self, op, tracers, params):
1542        vals_in, bdims_in = unzip2((t.val, t.vmap_dim) for t in tracers)
1543        val_outs, bdim_outs = op.vmap(self.dim_size, vals_in, bdims_in, **params)
1544        return [VMapTraceTensor(self, x, bd) for x, bd in list_zip(val_outs, bdim_outs)]
1545
1546    @staticmethod
1547    def move_vmap_dim(x, dim_size, src: int, dst: int):
1548        if src is None:  # unsqueeze and expand
1549            target_shape = list(x.shape)
1550            target_shape.insert(dst, dim_size)
1551            unsqueeze_shape = [1 if d == dst else target_shape[d] for d in range(len(target_shape))]
1552            x = x.reshape(tuple(unsqueeze_shape))
1553            x = x.expand(tuple(target_shape))
1554            return x
1555        elif src == dst:
1556            return x
1557        else:
1558            perm = [i for i in range(len(x.shape)) if i != src]
1559            perm.insert(dst, src)
1560            return x.permute(tuple(perm))
1561
1562
1563class JVPTraceTensor(TraceTensor):
1564    def __init__(self, trace, primal, tangent):
1565        self._trace = trace
1566        self.primal = primal
1567        self.tangent = tangent
1568
1569    @property
1570    def symval(self):
1571        return get_symval(self.primal)
1572
1573    @property
1574    def val(self):
1575        return self.primal
1576
1577    @property
1578    def dtype(self):
1579        return self.primal.dtype
1580
1581
1582class JVPTrace(Trace):
1583    def pure(self, val):
1584        if isinstance(val, PartialRunTrace):
1585            val = val.pval.const
1586        return JVPTraceTensor(self, val, backend.zeros_like(val))
1587
1588    def run_op(self, op, tracers, params):
1589        primals_in, tangents_in = unzip2((t.primal, t.tangent) for t in tracers)
1590        primals_out, tangents_out = op.jvp(primals_in, tangents_in, **params)
1591        return [JVPTraceTensor(self, x, t) for x, t in list_zip(primals_out, tangents_out)]
1592
1593
1594class ProgramTraceTensor(TraceTensor):
1595    __slots__ = ["symval"]
1596    symval: SymbolicTensor
1597
1598    def __init__(self, trace, symval):
1599        self._trace = trace
1600        self.symval = symval
1601
1602
1603class ProgramTrace(Trace):
1604    @property
1605    def builder(self):
1606        return self.main.global_data
1607
1608    def new_arg(self, symval) -> ProgramTraceTensor:
1609        symval = SymbolicTensor.like(symval)
1610        tracer = self.builder.new_tracer(self, symval)
1611        self.builder.tracer_to_var[id(tracer)] = Var(symval)
1612
1613        return tracer
1614
1615    def pure(self, val: Any) -> ProgramTraceTensor:
1616        # get_or_make_const_tracer
1617        tracer = self.builder.const_tracers.get(id(val))
1618        if tracer is None:
1619            tracer = self.builder.new_tracer(self, get_symval(val))
1620            self.builder.add_const(tracer, val)
1621        # print(self.builder.const_tracers)
1622        return tracer
1623
1624    def run_op(self, op, tracers, params):
1625        symvals_in = tree_map(lambda x: x.symval, tracers)
1626        symvals_out = op.typecheck(*symvals_in, **params)
1627
1628        out_tracers = [self.builder.new_tracer(self, a) for a in symvals_out]
1629        inputs = [self.builder.getvar(t) for t in tracers]
1630        outvars = [self.builder.add_var(t) for t in out_tracers]
1631
1632        self.builder.add_instruction(Instruction(op, inputs, params, outvars))
1633        return out_tracers
1634
1635
1636class ProgramBuilder:
1637    instructions: List[Instruction]
1638    tracer_to_var: Dict[int, Var]
1639    const_tracers: Dict[int, TraceTensor]
1640    constvals: Dict[Var, Any]
1641    tracers: List[ProgramTraceTensor]
1642
1643    def __init__(self):
1644        self.instructions = []
1645        self.tracer_to_var = {}
1646        self.const_tracers = {}
1647        self.constvals = {}
1648        self.tracers = []
1649
1650    def new_tracer(self, trace: ProgramTrace, symval: SymbolicTensor) -> ProgramTraceTensor:
1651        tracer = ProgramTraceTensor(trace, symval)
1652        self.tracers += [tracer]
1653        return tracer
1654
1655    def add_instruction(self, instruction: Instruction) -> None:
1656        self.instructions += [instruction]
1657
1658    def add_var(self, tracer: ProgramTraceTensor) -> Var:
1659        assert id(tracer) not in self.tracer_to_var
1660        var = self.tracer_to_var[id(tracer)] = Var(tracer.symval)
1661        return var
1662
1663    def getvar(self, tracer: ProgramTraceTensor) -> Var:
1664        var = self.tracer_to_var.get(id(tracer))
1665        assert var is not None
1666        return var
1667
1668    def add_const(self, tracer: ProgramTraceTensor, val: Any) -> Var:
1669        var = self.add_var(tracer)
1670        self.const_tracers[id(val)] = tracer
1671        self.constvals[var] = val
1672        return var
1673
1674    def build(self, in_tracers: Any, out_tracers: Any, static_args, name) -> Tuple[Program, List[Any]]:
1675        constvars, constvals = unzip2(self.constvals.items())
1676        t2v = lambda t: self.tracer_to_var[id(t)]
1677        in_binders = constvars + [t2v(t) for t in in_tracers]
1678        out_vars = [t2v(t) for t in out_tracers]
1679        program = Program(
1680            in_binders,
1681            self.instructions,
1682            out_vars,
1683            len(constvals),
1684            static_args,
1685            name,
1686        )
1687        typecheck_program(program)
1688        program, constvals = self._inline_literals(program, constvals)
1689        typecheck_program(program)
1690        # dblog(program, enable=backend.LOG_PROGRAM)
1691        return program, constvals
1692
1693    def _inline_literals(self, program: Program, consts: List[Any]) -> Tuple[Program, List[Any]]:
1694        const_binders, other_binders = split_list(program.in_binders, len(consts))
1695        scalars = [type(x) in TraceTensor.PYTHON_TYPES and not get_symval(x).shape for x in consts]
1696        new_const_binders, lit_binders = partition_list(scalars, const_binders)
1697        new_consts, lit_vals = partition_list(scalars, consts)
1698        literals = dict(list_zip(lit_binders, list_map(Lit, lit_vals)))
1699        new_outs = [literals.get(x, x) for x in program.outs]
1700        new_instructions = [
1701            Instruction(
1702                instruction.op,
1703                [literals.get(x, x) for x in instruction.inputs],
1704                instruction.params,
1705                instruction.out_binders,
1706            )
1707            for instruction in program.instructions
1708        ]
1709        new_program = Program(
1710            new_const_binders + other_binders,
1711            new_instructions,
1712            new_outs,
1713            len(new_consts),
1714            program.static_args,
1715            program.name,
1716        )
1717        return new_program, tuple(new_consts)
1718
1719    def get_current_scope_info(self):
1720        current_frame = inspect.currentframe()
1721        current_function_name = current_frame.f_code.co_name
1722        current_module_name = inspect.getmodulename(current_frame.f_code.co_filename)
1723        current_class_name = None
1724        for frame_info in inspect.getouterframes(current_frame):
1725            print(frame_info)
1726            frame_locals = frame_info.frame.f_locals
1727            print(frame_locals)
1728            if "self" in frame_locals:
1729                current_class_name = frame_locals["self"].__class__.__name__
1730                break
1731        return {
1732            "Function": current_function_name,
1733            "Module": current_module_name,
1734            "Class": current_class_name,
1735        }
1736
1737
1738class UndefinedPrimal(NamedTuple):
1739    symval: SymbolicTensor
1740
1741    @property
1742    def shape(self):
1743        return self.symval.shape
1744
1745    @property
1746    def dtype(self):
1747        return self.symval.dtype
1748
1749    @property
1750    def device(self):
1751        return self.symval.device
1752
1753    @property
1754    def ndim(self):
1755        return self.symval.ndim
1756
1757    def __repr__(self):
1758        return f"<UndefinedPrimal: symval={self.symval}>"
1759
1760    str_short = __repr__
1761
1762
1763class PartialValue(NamedTuple):
1764    symval: SymbolicTensor
1765    const: Optional[Any]
1766
1767    is_known = property(lambda self: self.const is not None)
1768    is_unknown = property(lambda self: self.const is None)
1769
1770
1771class LambdaBindingDraft(NamedTuple):
1772    pass
1773
1774
1775class ConstDraft(NamedTuple):
1776    val: Any
1777
1778
1779class InstructionDraft(NamedTuple):
1780    prim: Operator
1781    tracers_in: List["PartialRunTraceTensor"]
1782    params: Dict[str, Any]
1783    symvals_out: List[SymbolicTensor]
1784    tracer_refs_out: List[weakref.ReferenceType["PartialRunTraceTensor"]]
1785
1786
1787ProgramDraft = Union[LambdaBindingDraft, ConstDraft, InstructionDraft]
1788
1789
1790class PartialRunTraceTensor(TraceTensor):
1791    def __init__(self, trace, pval, draft):
1792        self._trace = trace
1793        self.pval = pval
1794        self.draft = draft
1795
1796    symval = property(lambda self: self.pval.symval)
1797    val = property(lambda self: self.pval.const)
1798
1799    def full_lower(self):
1800        if self.pval.is_known:
1801            return full_lower(self.pval.const)
1802        return self
1803
1804
1805class PartialRunTrace(Trace):
1806    def new_arg(self, pval: PartialValue) -> Any:
1807        return PartialRunTraceTensor(self, pval, LambdaBindingDraft())
1808
1809    def pure(self, val: Any) -> PartialRunTraceTensor:
1810        return PartialRunTraceTensor(self, make_known_pval(val), None)
1811
1812    def instantiate_const(self, tracer: PartialRunTraceTensor) -> PartialRunTraceTensor:
1813        if tracer.pval.is_unknown:
1814            return tracer
1815        else:
1816            pval = make_unknown_pval(SymbolicTensor.like(tracer.symval))
1817            return PartialRunTraceTensor(self, pval, ConstDraft(tracer.pval.const))
1818
1819    def run_op(self, op, tracers, params):
1820        is_knowns = tuple(t.pval.is_known for t in tracers)
1821
1822        if all(is_knowns):
1823            return bind(op, *list_map(full_lower, tracers), **params)
1824        return op.partial_run(self, tracers, **params)
1825
1826
1827trace_stack: List[MainTrace] = []
1828stashed_trace: Optional[MainTrace] = None
1829trace_stack += [MainTrace(0, RunTrace, None)]
1830
1831
1832class UndefBackend:
1833    def __getattr__(self, attr):
1834        raise NotImplementedError("Backend not init yet with slope.core.set_backend(backend)")
1835
1836
1837backend = UndefBackend()
1838
1839
1840def set_backend(name, where="slope.backends"):
1841    global backend
1842    backend = importlib.import_module(f"{where}.{name}").backend
1843    import slope.nn as nn
1844
1845    # backend.register_node(nn.Module, nn.Module.flatten, nn.Module.unflatten, "Module")
1846
1847    dblog(f"slope backend is {backend}", enable=backend.LOG_INIT)
1848
1849
1850def stack_str():
1851    ret = ""
1852    for trace in trace_stack:
1853        ret += f"{trace.level}: {trace.trace_type.__name__}\t{trace.global_data=}\n"
1854    return ret
1855
1856
1857def make_known_pval(val: Any):
1858    return PartialValue(get_symval(val), val)
1859
1860
1861def make_unknown_pval(symval: SymbolicTensor):
1862    return PartialValue(symval, None)
1863
1864
1865def get_symval(x):
1866    if isinstance(x, TraceTensor):
1867        return x.symval
1868    elif type(x) in TraceTensor.PYTHON_TYPES:
1869        return backend.tensor(x)
1870    elif isinstance(x, Tensor):
1871        return x
1872    elif isinstance(x, SymbolicTensor):
1873        return x
1874    else:
1875        raise TypeError(type(x))
1876
1877
1878def tree_flatten(x: Any) -> Any:
1879    def _tree_flatten(x_: Any) -> Tuple[Iterable, Union[TreeDef, Leaf]]:
1880        node_type = None
1881        for k in backend.node_types.keys():
1882            if isinstance(x_, k):
1883                node_type = backend.node_types[k]
1884
1885        if node_type is not None:
1886            node_metadata, children = node_type.flatten(x_)
1887            children_flat, child_trees = unzip2(list_map(_tree_flatten, children))
1888            children_iter = itertools.chain.from_iterable(children_flat)
1889            treedef = TreeDef(node_type, node_metadata, tuple(child_trees))
1890            return children_iter, treedef
1891        else:
1892            return (x_,), Leaf(x_)
1893
1894    children_iter, treedef = _tree_flatten(x)
1895    return tuple(children_iter), treedef
1896
1897
1898def tree_unflatten(treedef: TreeDef, xs: Tuple[Any]) -> Any:
1899    def _tree_unflatten(treedef_: TreeDef, xs_: Iterator) -> Any:
1900        if isinstance(treedef_, Leaf):
1901            dblog(f"    tree leaf found: {xs_}\n", enable=backend.LOG_TREE)
1902            return next(xs_)
1903        else:
1904            dblog(f"    now\n  {treedef_}", enable=backend.LOG_TREE)
1905            children = (_tree_unflatten(t, xs_) for t in treedef_.child_treedefs)
1906            dblog(f"{children=}\n", enable=backend.LOG_TREE)
1907            return treedef_.node_type.unflatten(treedef_.node_metadata, children)
1908
1909    dblog(f"unflattening {treedef}", enable=backend.LOG_TREE)
1910    return _tree_unflatten(treedef, iter(xs))
1911    # with Timing(f"\nTREE:\n{treedef}"):
1912    #     ret = _tree_unflatten(treedef, iter(xs))
1913    # return ret
1914
1915
1916def tree_transpose(
1917    outer_treedef: TreeDef,
1918    inner_treedef: TreeDef,
1919    tree_to_transpose: Any,
1920) -> Any:
1921    flat, treedef = tree_flatten(tree_to_transpose)
1922    inner_size = inner_treedef.num_leaves
1923    outer_size = outer_treedef.num_leaves
1924    if treedef.num_leaves != (inner_size * outer_size):
1925        raise TypeError
1926    iter_flat = iter(flat)
1927    lol = [[next(iter_flat) for _ in range(inner_size)] for __ in range(outer_size)]
1928    permuted_lol = zip(*lol)
1929    subtrees = map(partial(tree_unflatten, outer_treedef), permuted_lol)
1930    return tree_unflatten(inner_treedef, subtrees)
1931
1932
1933def flatten_fn(f, in_tree, *, has_aux=False):
1934    store = Store()
1935
1936    def flat_fn(*args_flat, **params):
1937        tree_args = tree_unflatten(in_tree, args_flat)
1938        out = f(*tree_args, **params)
1939        if has_aux:
1940            out, aux = out
1941        out_flat, out_tree = tree_flatten(out)
1942        store.set_value(out_tree)
1943        return (out_flat, aux) if has_aux else out_flat
1944
1945    return flat_fn, store
1946
1947
1948def tree_map(f: Callable[..., Any], tree, *rest, out_leaf=False) -> Any:
1949    leaves, treedef = tree_flatten(tree)
1950    if len(rest) == 0:
1951        out_tree_flat = tuple(f(leaf) for leaf in leaves)
1952        out_tree = tree_unflatten(treedef, out_tree_flat)
1953    else:
1954        all_leaves = [leaves]
1955        for t in rest:
1956            t_leaves, t_treedef = tree_flatten(t)
1957            assert t_treedef == treedef
1958            all_leaves += [t_leaves]
1959
1960        out_tree_flat = tuple(f(*xs) for xs in zip(*all_leaves))
1961        out_tree = tree_unflatten(treedef, out_tree_flat)
1962    ret = out_tree
1963    if out_leaf:
1964        ret = (ret, tree_flatten(out_tree_flat[0]))
1965    return ret
1966
1967
1968@contextmanager
1969def new_main_trace(trace_type: Type["Trace"], global_data=None):
1970    global trace_stack
1971    level = len(trace_stack)
1972    main = MainTrace(level, trace_type, global_data)
1973    trace_stack += [main]
1974
1975    try:
1976        yield main
1977    finally:
1978        trace_stack.pop()
1979
1980
1981def bind(op, *args, **params):
1982    top_trace = find_top_trace(args)
1983    tracers = tuple([full_raise(top_trace, arg) for arg in args])
1984    outs = top_trace.run_op(op, tracers, params)
1985    lowered = tuple([full_lower(out) for out in outs])
1986    return lowered
1987
1988
1989def find_top_trace(xs) -> Trace:
1990    arrs = []
1991
1992    def get_arr_from_seq(seq):
1993        nonlocal arrs
1994        for x in seq:
1995            if type(x) in (tuple, list):
1996                get_arr_from_seq(x)
1997            elif isinstance(x, TraceTensor):
1998                arrs += [x]
1999
2000    get_arr_from_seq(xs)
2001    arrs = tuple(arrs)
2002    top_main = max(
2003        (x._trace.main for x in arrs),
2004        default=trace_stack[0],
2005        key=operator_py.attrgetter("level"),
2006    )
2007    if stashed_trace and stashed_trace.level > top_main.level:
2008        top_main = stashed_trace
2009    return top_main.trace_type(top_main)
2010
2011
2012def full_raise(trace: Trace, val: Any) -> TraceTensor:
2013    if not isinstance(val, TraceTensor):
2014        return trace.pure(val)
2015    level = trace.main.level
2016    if val._trace.main is trace.main:
2017        return val
2018    elif val._trace.main.level < level:
2019        return trace.pure(val)
2020    elif val._trace.main.level > level:
2021        raise Exception(f"Can't lift level {val._trace.main.level} to {level}.")
2022    else:
2023        raise Exception(f"Different traces at same level: {val._trace}, {trace}.")
2024
2025
2026def full_lower(val: Any):
2027    if isinstance(val, TraceTensor):
2028        return val.full_lower()
2029    elif type(val) in (list, tuple):
2030        return tuple(full_lower(v) for v in val)
2031    else:
2032        return val
2033
2034
2035def typecheck_program(program: Program) -> ProgramType:
2036    env: Set[Var] = set()
2037
2038    for v in program.in_binders:
2039        if v in env:
2040            raise TypeError
2041        env.add(v)
2042
2043    for instruction in program.instructions:
2044        in_types = [typecheck_atom(env, x) for x in instruction.inputs]
2045        out_types = instruction.op.typecheck(*in_types, **instruction.params)
2046        for out_binder, out_type in list_zip(instruction.out_binders, out_types):
2047            if not out_type == out_binder.symval:
2048                raise TypeError
2049        for out_binder in instruction.out_binders:
2050            if out_binder in env:
2051                raise TypeError
2052            env.add(out_binder)
2053
2054    in_types = [v.symval for v in program.in_binders]
2055    out_types = [typecheck_atom(env, x) for x in program.outs]
2056    return ProgramType(tuple(in_types), tuple(out_types))
2057
2058
2059def typecheck_atom(env: Set[Var], x: Atom) -> SymbolicTensor:
2060    if isinstance(x, Var):
2061        if x not in env:
2062            raise TypeError("unbound variable")
2063        return x.symval
2064    elif isinstance(x, Lit):
2065        return get_symval(x.val)
2066    else:
2067        assert False
2068
2069
2070def run_program(program: Program, args: List[Any]) -> List[Any]:
2071    env: Dict[Var, Any] = {}
2072
2073    def read(x: Atom) -> Any:
2074        return env[x] if type(x) is Var else x.val
2075
2076    def write(v: Var, val: Any) -> None:
2077        assert v not in env  # single-assignment
2078        env[v] = val
2079
2080    list_map(write, program.in_binders, args)
2081    for instruction in program.instructions:
2082        in_vals = list_map(read, instruction.inputs)
2083        outs = bind(instruction.op, *in_vals, **instruction.params)
2084        list_map(write, instruction.out_binders, outs)
2085    return list_map(read, program.outs)
2086
2087
2088def program_as_fun(program: Program):
2089    return lambda *args: run_program(program, args)
2090
2091
2092def vmap_flat(f, in_dim, out_dim, dim_size, *args):
2093    if dim_size is None:
2094        dims = set([x.shape[d] for x, d in list_zip(args, in_dim) if d is not None])
2095        assert len(dims) == 1
2096        (dim_size,) = dims
2097    with new_main_trace(VMapTrace, dim_size) as main:
2098        trace = VMapTrace(main)
2099        tracers_in = [VMapTraceTensor(trace, x, dim) if dim is not None else x for x, dim in list_zip(args, in_dim)]
2100        outs = f(*tracers_in)
2101        tracers_out = [full_raise(trace, out) for out in outs]
2102        vals_out, y_vmap_dims = unzip2((t.val, t.vmap_dim) for t in tracers_out)
2103    ret = [VMapTrace.move_vmap_dim(val_out, dim_size, bdim, out_dim) for val_out, bdim, out_dim in zip(vals_out, y_vmap_dims, out_dim)]
2104    return ret
2105
2106
2107def vmap(f, in_dim=0, out_dim=0, dim_size=None):
2108    def batched_f(*args):
2109        nonlocal in_dim, out_dim, dim_size
2110        args_flat, in_tree = tree_flatten(args)
2111        in_dim = (in_dim,) * len(args) if isinstance(in_dim, int) else in_dim
2112        out_dim = (out_dim,) * len(args) if isinstance(out_dim, int) else out_dim
2113        in_dim_flat, in_dim_tree = tree_flatten(in_dim)
2114        out_dim_flat, out_dim_tree = tree_flatten(out_dim)
2115        if not (in_tree == in_dim_tree == out_dim_tree):
2116            raise TypeError(f"\n{in_tree}\n!=\n{in_dim_tree}!=\n{out_dim_tree}")
2117        f_flat, out_tree_store = flatten_fn(f, in_tree)
2118        # if len(args_flat) > len(in_dim_flat):
2119        #     in_dim_flat = (in_dim[0],) * len(args_flat)
2120        outs_flat = vmap_flat(f_flat, in_dim_flat, out_dim_flat, dim_size, *args_flat)
2121        return tree_unflatten(out_tree_store(), outs_flat)
2122
2123    return batched_f
2124
2125
2126def jvp_flat(f, primals, tangents, *, has_aux, global_data, **static_args):
2127    with new_main_trace(JVPTrace, global_data) as main:
2128        trace = JVPTrace(main)
2129        tracers_in = [JVPTraceTensor(trace, x, t) for x, t in list_zip(primals, tangents)]
2130        jvp_flat_ret = f(*tracers_in, **static_args)
2131        if has_aux:
2132            (outs, aux) = jvp_flat_ret
2133            # aux_ = aux
2134            aux = tree_map(lambda x: x.primal, aux)
2135            # aux = tree_map(lambda x: x.full_lower(), aux)
2136            #
2137        else:
2138            outs = jvp_flat_ret
2139        tracers_out = [full_raise(trace, out) for out in outs]
2140        primals_out, tangents_out = unzip2((t.primal, t.tangent) for t in tracers_out)
2141    return ((primals_out, tangents_out), aux) if has_aux else (primals_out, tangents_out)
2142
2143
2144def jvp(f, primals, tangents, *, has_aux=False, global_data=None, **static_args):
2145    primals_flat, in_tree = tree_flatten(primals)
2146    tangents_flat, in_tree2 = tree_flatten(tangents)
2147    for p, t in zip(primals_flat, tangents_flat):
2148        assert p.shape == t.shape, f"{p.shape=} != {t.shape=}"
2149        assert p.dtype == t.dtype, f"{p.dtype=} != {t.dtype=}"
2150        assert p.device == t.device, f"{p.device=} != {t.device=}"
2151    if in_tree != in_tree2:
2152        raise TypeError
2153    f, out_tree_store = flatten_fn(f, in_tree, has_aux=has_aux)
2154    jvp_ret = jvp_flat(
2155        f,
2156        primals_flat,
2157        tangents_flat,
2158        has_aux=has_aux,
2159        global_data=global_data,
2160        **static_args,
2161    )
2162    if has_aux:
2163        (primals_out_flat, tangents_out_flat), aux = jvp_ret
2164    else:
2165        (primals_out_flat, tangents_out_flat) = jvp_ret
2166    primals_out = tree_unflatten(out_tree_store(), primals_out_flat)
2167    tangents_out = tree_unflatten(out_tree_store(), tangents_out_flat)
2168    return ((primals_out, tangents_out), aux) if has_aux else (primals_out, tangents_out)
2169
2170
2171def jacfwd(f, argnums=0, has_aux=False):
2172    def jvp_fn(x):
2173        return jvp(f, x, (backend.eye(len(x)),), has_aux=has_aux)
2174    return vmap(jvp_fn, in_dim=argnums)
2175
2176def jacrev(f, argnums=0, has_aux=False):
2177    def grad_f(x):
2178        return grad(lambda x: f(x) @ backend.eye(f(x).shape[0]), argnums=argnums, has_aux=has_aux)(x)
2179    return vmap(grad_f)
2180
2181## arange version
2182# def jacrev(f, x):
2183#     def grad_f_i(x, i):
2184#         return grad(lambda x: f(x)[i])(x)
2185#     return vmap(lambda i: grad_i(x, i))(backend.arange(f(x).shape[0]))
2186
2187def hessian(fn, argnums=0, has_aux=False):
2188    return jacrev(jacrev(fn, argnums=argnums, has_aux=has_aux))
2189
2190@contextmanager
2191def stash_trace(main: MainTrace):
2192    global stashed_trace
2193    prev_stashed_trace, stashed_trace = stashed_trace, main
2194    try:
2195        yield
2196    finally:
2197        stashed_trace = prev_stashed_trace
2198
2199
2200@contextmanager
2201def symbolic_run():
2202    global trace_stack
2203    level = len(trace_stack)
2204    main = MainTrace(level, SymbolicRunTrace, global_data=None)
2205    trace_stack += [main]
2206    global stashed_trace
2207    prev_stashed_trace, stashed_trace = stashed_trace, main
2208    try:
2209        yield
2210    finally:
2211        stashed_trace = prev_stashed_trace
2212        trace_stack.pop()
2213
2214
2215@lru_cache_verbose()
2216def make_program(f: Callable, *symvals_in: SymbolicTensor, static_args, name) -> Tuple[Program, List[Any], TreeDef]:
2217    symvals_in, in_tree = tree_flatten(symvals_in)
2218    f, out_tree_store = flatten_fn(f, in_tree)
2219    builder = ProgramBuilder()
2220    with new_main_trace(ProgramTrace, builder) as main:
2221        with stash_trace(main):
2222            trace = ProgramTrace(main)
2223            tracers_in = [trace.new_arg(symval) for symval in symvals_in]
2224            outs = f(*tracers_in, **{k: v for k, v in static_args})
2225            tracers_out = [full_raise(trace, out) if isinstance(out, ProgramTraceTensor) else out.val for out in outs]
2226            program, consts = builder.build(tracers_in, tracers_out, static_args, name)
2227
2228    return program, consts, out_tree_store()
2229
2230
2231@lru_cache_verbose()
2232def vmap_program(program: Program, dim_size, dims_in) -> tuple[Program, list[Any]]:
2233    def unmapped_symval(axis_size: int, batch_dim, symval: SymbolicTensor) -> SymbolicTensor:
2234        if batch_dim is None:
2235            return symval
2236        else:
2237            shape = list(symval.shape)
2238            shape.insert(batch_dim, axis_size)
2239            return symval.like(shape=tuple(shape))
2240
2241    vmap_traceable = vmap(program_as_fun(program), tuple(dims_in))
2242    in_symvals = [unmapped_symval(dim_size, d, v.symval) for v, d in zip(program.in_binders, dims_in)]
2243    program, consts, _ = make_program(
2244        vmap_traceable,
2245        *in_symvals,
2246        static_args=program.static_args,
2247        name=f"vmap_{program.name}",
2248    )
2249    return program, consts
2250
2251
2252@lru_cache_verbose()
2253def jvp_program(program: Program) -> Tuple[Program, List[Any]]:
2254    def jvp_traceable(*primals_and_tangents):
2255        n = len(primals_and_tangents) // 2
2256        primals, tangents = primals_and_tangents[:n], primals_and_tangents[n:]
2257        return jvp(program_as_fun(program), primals, tangents)
2258
2259    in_symvals = tree_map(lambda v: v.symval, program.in_binders)
2260    new_program, new_consts, _ = make_program(
2261        jvp_traceable,
2262        *in_symvals,
2263        *in_symvals,
2264        static_args=program.static_args,
2265        name=f"{program.name}_jvp",
2266    )
2267    return new_program, new_consts
2268
2269
2270def partial_run_flat(
2271    f: Callable, pvals_in: List["PartialValue"], has_aux, global_data=None
2272) -> Tuple[Program, List["PartialValue"], List[Any]]:
2273    with new_main_trace(PartialRunTrace, global_data) as main:
2274        trace = PartialRunTrace(main)
2275        tracers_in = [trace.new_arg(pval) for pval in pvals_in]
2276        outs = f(*tracers_in)
2277        if has_aux:
2278            outs, aux = outs
2279        tracers_out = [full_raise(trace, out) for out in outs]
2280        pvals_out = [t.pval for t in tracers_out]
2281        unk_tracers_in = [t for t in tracers_in if t.pval.is_unknown]
2282        unk_tracers_out = [t for t in tracers_out if t.pval.is_unknown]
2283        program, consts = tracers_to_program(unk_tracers_in, unk_tracers_out)
2284
2285    return (program, pvals_out, consts, aux) if has_aux else (program, pvals_out, consts)
2286
2287
2288def partial_run_program(
2289    program: Program,
2290    in_unknowns: List[bool],
2291    instantiate: Optional[List[bool]] = None,
2292) -> Tuple[Program, Program, List[bool], int]:
2293    env: Dict[Var, bool] = {}
2294    residuals: Set[Var] = set()
2295
2296    def read(x: Atom) -> bool:
2297        return type(x) is Var and env[x]
2298
2299    def write(unk: bool, v: Var) -> None:
2300        env[v] = unk
2301
2302    instructions1, instructions2 = [], []
2303    list_map(write, in_unknowns, program.in_binders)
2304
2305    for instruction in program.instructions:
2306        unks_in = list_map(read, instruction.inputs)
2307        (
2308            instruction1,
2309            instruction2,
2310            unks_out,
2311            res,
2312        ) = instruction.op.partial_run_instruction(unks_in, instruction)
2313        if instruction1 is not None:
2314            instructions1 += [instruction1]
2315        if instruction2 is not None:
2316            instructions2 += [instruction2]
2317        if res is not None:
2318            residuals.update(res)
2319        list_map(write, unks_out, instruction.out_binders)
2320
2321    out_unknowns = list_map(read, program.outs)
2322    if instantiate is not None:
2323        for v, uk, inst in zip(program.outs, out_unknowns, instantiate):
2324            if inst and not uk:
2325                if type(v) is Var:
2326                    residuals.add(v)
2327        out_unknowns = list_map(operator_py.or_, out_unknowns, instantiate)
2328
2329    residuals, num_res = list(residuals), len(residuals)
2330    assert all(type(v) is Var for v in residuals), residuals
2331
2332    ins1, ins2 = partition_list(in_unknowns, program.in_binders)
2333    outs1, outs2 = partition_list(out_unknowns, program.outs)
2334
2335    program1 = Program(
2336        ins1,
2337        instructions1,
2338        outs1 + residuals,
2339        0,
2340        program.static_args,
2341        f"{program.name}_partial1",
2342    )
2343    program2 = Program(
2344        residuals + ins2,
2345        instructions2,
2346        outs2,
2347        0,
2348        program.static_args,
2349        f"{program.name}_partial2",
2350    )
2351    typecheck_partial_run_program(program, in_unknowns, out_unknowns, program1, program2)
2352
2353    return program1, program2, out_unknowns, num_res
2354
2355
2356def typecheck_partial_run_program(program, in_unknowns, out_unknowns, program1, program2):
2357    programty = typecheck_program(program)  # (a1,  a2) -> (b1, b2 )
2358    program1ty = typecheck_program(program1)  #  a1       -> (b1, res)
2359    program2ty = typecheck_program(program2)  # (res, a2) -> b2
2360
2361    a1, a2 = partition_list(in_unknowns, programty.in_types)
2362    b1, b2 = partition_list(out_unknowns, programty.out_types)
2363    b1_, res = split_list(program1ty.out_types, len(b1))
2364    res_, a2_ = split_list(program2ty.in_types, len(res))
2365    b2_ = program2ty.out_types
2366
2367    a1 = tuple(a1)
2368    a2, a2_ = tuple(a2), tuple(a2_)
2369    b1, b1_ = tuple(b1), tuple(b1_)
2370    b2, b2_ = tuple(b2), tuple(b2_)
2371    res, res_ = tuple(res), tuple(res_)
2372
2373    if program1ty.in_types != a1:
2374        raise TypeError
2375    if program2ty.out_types != b2:
2376        raise TypeError
2377    if b1 != b1_:
2378        raise TypeError
2379    if res != res_:
2380        raise TypeError
2381    if a2 != a2_:
2382        raise TypeError
2383    if b2 != b2_:
2384        raise TypeError
2385
2386
2387def linearize_flat(f, *primals_in, has_aux):
2388    pvals_in = [make_known_pval(x) for x in primals_in] + [make_unknown_pval(SymbolicTensor.like(get_symval(x))) for x in primals_in]
2389
2390    def f_jvp(*primals_tangents_in):
2391        jvp_ret = jvp(f, *split_half(primals_tangents_in), has_aux=has_aux)
2392        if has_aux:
2393            (primals_out, tangents_out), aux = jvp_ret
2394            return ((*primals_out, *tangents_out), aux)
2395        else:
2396            primals_out, tangents_out = jvp_ret
2397            return (*primals_out, *tangents_out)
2398
2399    partial_run_flat_ret = partial_run_flat(f_jvp, pvals_in, has_aux)
2400    if has_aux:
2401        program, pvals_out, consts, aux = partial_run_flat_ret
2402    else:
2403        program, pvals_out, consts = partial_run_flat_ret
2404    primal_pvals, _ = split_half(pvals_out)
2405    assert all(pval.is_known for pval in primal_pvals)
2406    primals_out = [pval.const for pval in primal_pvals]
2407    f_lin = lambda *tangents: run_program(program, [*consts, *tangents])
2408    return (primals_out, f_lin, aux) if has_aux else (primals_out, f_lin)
2409
2410
2411def linearize(f, *primals_in, has_aux=False):
2412    primals_in_flat, in_tree = tree_flatten(primals_in)
2413    f, out_tree_store = flatten_fn(f, in_tree, has_aux=has_aux)
2414    linearize_flat_ret = linearize_flat(f, *primals_in_flat, has_aux=has_aux)
2415    if has_aux:
2416        primals_out_flat, f_lin_flat, aux = linearize_flat_ret
2417    else:
2418        primals_out_flat, f_lin_flat = linearize_flat_ret
2419
2420    primals_out = tree_unflatten(out_tree_store(), primals_out_flat)
2421
2422    def f_lin(*tangents_in):
2423        tangents_in_flat, in_tree2 = tree_flatten(tangents_in)
2424        if in_tree != in_tree2:
2425            raise TypeError
2426        tangents_out_flat = f_lin_flat(*tangents_in_flat)
2427        return tree_unflatten(out_tree_store(), tangents_out_flat)
2428
2429    return (primals_out, f_lin, aux) if has_aux else (primals_out, f_lin)
2430
2431
2432def tracers_to_program(
2433    tracers_in: List["PartialRunTraceTensor"],
2434    tracers_out: List["PartialRunTraceTensor"],
2435):
2436    def tracer_parents(t: PartialRunTraceTensor) -> List[PartialRunTraceTensor]:
2437        return t.draft.tracers_in if isinstance(t.draft, InstructionDraft) else []
2438
2439    def draft_to_instruction(tracer_to_var: Dict[int, Var], draft: InstructionDraft) -> Instruction:
2440        inputs = [tracer_to_var[id(t)] for t in draft.tracers_in]
2441        out_binders = [Var(symval) for symval in draft.symvals_out]
2442        for t_ref, var in list_zip(draft.tracer_refs_out, out_binders):
2443            if t_ref() is not None:
2444                tracer_to_var[id(t_ref())] = var
2445        return Instruction(draft.prim, inputs, draft.params, out_binders)
2446
2447    tracer_to_var: Dict[int, Var] = {id(t): Var(SymbolicTensor.like(t.symval)) for t in tracers_in}
2448    constvar_to_val: Dict[int, Any] = {}
2449    constid_to_var: Dict[int, Var] = {}
2450    processed_instructions: Set[int] = set()
2451    instructions: List[Instruction] = []
2452    for t in toposort(tracers_out, tracer_parents):
2453        if isinstance(t.draft, LambdaBindingDraft):
2454            assert id(t) in set(list_map(id, tracers_in))
2455        elif isinstance(t.draft, ConstDraft):
2456            val = t.draft.val
2457            var = constid_to_var.get(id(val))
2458            if var is None:
2459                symval = SymbolicTensor.like(get_symval(val))
2460                var = constid_to_var[id(val)] = Var(symval)
2461                constvar_to_val[var] = val
2462            tracer_to_var[id(t)] = var
2463        elif isinstance(t.draft, InstructionDraft):
2464            if id(t.draft) not in processed_instructions:
2465                instructions += [draft_to_instruction(tracer_to_var, t.draft)]
2466                processed_instructions.add(id(t.draft))
2467        else:
2468            raise TypeError(t.draft)
2469
2470    constvars, constvals = unzip2(constvar_to_val.items())
2471    in_binders = constvars + [tracer_to_var[id(t)] for t in tracers_in]
2472    # out_vars = [tracer_to_var[id(t)] for t in tracers_out if id(t) in tracer_to_var]
2473    out_vars = [tracer_to_var[id(t)] for t in tracers_out]
2474    program = Program(tuple(in_binders), tuple(instructions), tuple(out_vars))
2475    typecheck_program(program)
2476    return program, constvals
2477
2478
2479def toposort(out_nodes: List[Any], parents: Callable[[Any], List[Any]]):
2480    def check_toposort(nodes: List[Any], parents: Callable[[Any], List[Any]]):
2481        seen = set()
2482        for node in nodes:
2483            assert all(id(parent) in seen for parent in parents(node))
2484            seen.add(id(node))
2485
2486    def remove_duplicates(lst):
2487        seen = set()
2488        return [x for x in lst if id(x) not in seen and not seen.add(id(x))]
2489
2490    if not out_nodes:
2491        return []
2492    out_nodes = remove_duplicates(out_nodes)
2493
2494    child_counts = {}
2495    stack = list(out_nodes)
2496    while stack:
2497        node = stack.pop()
2498        if id(node) in child_counts:
2499            child_counts[id(node)] += 1
2500        else:
2501            child_counts[id(node)] = 1
2502            stack.extend(parents(node))
2503    for node in out_nodes:
2504        child_counts[id(node)] -= 1
2505
2506    sorted_nodes = []
2507    childless_nodes = [node for node in out_nodes if not child_counts[id(node)]]
2508    while childless_nodes:
2509        node = childless_nodes.pop()
2510        sorted_nodes += [node]
2511        for parent in parents(node):
2512            if child_counts[id(parent)] == 1:
2513                childless_nodes += [parent]
2514            else:
2515                child_counts[id(parent)] -= 1
2516
2517    sorted_nodes = sorted_nodes[::-1]
2518    check_toposort(sorted_nodes, parents)
2519    return sorted_nodes
2520
2521
2522def vjp_flat(f, *primals_in, has_aux=False, **static_args):
2523    pvals_in = [make_known_pval(x) for x in primals_in] + [make_unknown_pval(SymbolicTensor.like(get_symval(x))) for x in primals_in]
2524    primal_pvals_in, tangent_pvals_in = split_half(pvals_in)
2525    del primal_pvals_in
2526
2527    def f_jvp(*primals_tangents_in):
2528        jvp_ret = jvp(
2529            f,
2530            *split_half(primals_tangents_in),
2531            has_aux=has_aux,
2532            global_data="vjp",
2533            **static_args,
2534        )
2535        if has_aux:
2536            ((primals_out, tangents_out), aux) = jvp_ret
2537        else:
2538            (primals_out, tangents_out) = jvp_ret
2539        return ([*primals_out, *tangents_out], aux) if has_aux else [*primals_out, *tangents_out]
2540
2541    partial_run_flat_ret = partial_run_flat(f_jvp, pvals_in, has_aux, "vjp")
2542    if has_aux:
2543        program, pvals_out, consts, aux = partial_run_flat_ret
2544    else:
2545        program, pvals_out, consts = partial_run_flat_ret
2546
2547    primal_pvals, tangent_pvals = split_half(pvals_out)
2548    del tangent_pvals
2549    assert all(pval.is_known for pval in primal_pvals)
2550    primals_out_flat = [pval.const for pval in primal_pvals]
2551    transpose_inputs = consts + [UndefinedPrimal(t.symval) for t in tangent_pvals_in]
2552
2553    def f_vjp_flat(*cotangents):
2554        # return backward_pass(program, transpose_inputs, cotangents)
2555        undef_primals = tuple(isinstance(x, UndefinedPrimal) for x in transpose_inputs)
2556        transposed_program, new_consts = transpose_program(program, undef_primals)
2557        residuals, _ = partition_list(undef_primals, transpose_inputs)
2558        outs = run_program(transposed_program, (*new_consts, *residuals, *cotangents))
2559        return outs
2560
2561    return (primals_out_flat, f_vjp_flat, aux) if has_aux else (primals_out_flat, f_vjp_flat)
2562
2563
2564def vjp(f, *primals_in, has_aux=False, **static_args):
2565    primals_in_flat, in_tree = tree_flatten(primals_in)
2566    f, out_tree_store = flatten_fn(f, in_tree, has_aux=has_aux)
2567    vjp_ret = vjp_flat(f, *primals_in_flat, has_aux=has_aux, **static_args)
2568    if has_aux:
2569        primals_out_flat, f_vjp_flat, aux = vjp_ret
2570    else:
2571        primals_out_flat, f_vjp_flat = vjp_ret
2572    primals_out = tree_unflatten(out_tree_store(), primals_out_flat)
2573
2574    def f_vjp(*cotangents_out):
2575        cotangents_out_flat, _ = tree_flatten(cotangents_out)
2576        cotangents_in_flat = f_vjp_flat(*cotangents_out_flat)
2577
2578        return tree_unflatten(in_tree, cotangents_in_flat)
2579
2580    return (primals_out, f_vjp, aux) if has_aux else (primals_out, f_vjp)
2581
2582
2583NullCotangent = None
2584
2585
2586def backward_pass(program: Program, args: List[Any], cotangents: List[Any]) -> List[Any]:
2587    primal_env: Dict[Var, Any] = {}
2588    ct_env: Dict[Var, Any] = {}
2589
2590    def read_primal(x: Atom) -> Any:
2591        return primal_env.get(x, UndefinedPrimal(x.symval)) if type(x) is Var else x.val
2592
2593    def write_primal(v: Var, val: Any) -> None:
2594        if type(val) is not UndefinedPrimal:
2595            primal_env[v] = val
2596
2597    def read_cotangent(v: Var) -> Any:
2598        return ct_env.pop(v, backend.zeros(v.symval.shape, v.symval.dtype))
2599
2600    def write_cotangent(x: Atom, ct: Any):
2601        if type(x) is Var and ct is not NullCotangent:
2602            ct_env[x] = (ct_env[x] + ct) if x in ct_env else ct
2603
2604    list_map(write_primal, program.in_binders, args)
2605    list_map(write_cotangent, program.outs, cotangents)
2606    for instruction in program.instructions[::-1]:
2607        primals_in = list_map(read_primal, instruction.inputs)
2608        cotangents_in = list_map(read_cotangent, instruction.out_binders)
2609        inp, params = primals_in, instruction.params
2610        cotangents_out = instruction.op.T(cotangents_in, *inp, **params)
2611        list_map(write_cotangent, instruction.inputs, cotangents_out)
2612
2613    ret = [read_cotangent(v) for v, x in list_zip(program.in_binders, args) if isinstance(x, UndefinedPrimal)]
2614    return ret
2615
2616
2617@lru_cache_verbose()
2618def transpose_program(program: Program, undef_primals: tuple[bool, ...]) -> tuple[Program, list[Any]]:
2619    symvals_in, symvals_out = typecheck_program(program)
2620    traceable = partial(backward_pass, program)
2621    ()
2622    args = [UndefinedPrimal(a) if u else a for a, u in zip(symvals_in, undef_primals)]
2623    trans_program, consts, _ = make_program(
2624        traceable,
2625        tuple(args),
2626        tuple(symvals_out),
2627        static_args=program.static_args,
2628        name=f"{program.name}_T",
2629    )
2630    typecheck_program(trans_program)
2631
2632    return trans_program, consts
2633
2634
2635def grad(f, argnums=(0,), argnames="", has_aux=False, return_value=False):
2636    f, rejit = (f, False) if not isinstance(f, jit) else (f.f, True)
2637    if isinstance(argnums, int):
2638        argnums = (argnums,)
2639
2640    def gfn(x, *xs, **static_args):
2641        vjp_ret = vjp(f, x, *xs, has_aux=has_aux, **static_args)
2642        if has_aux:
2643            y, f_vjp, aux = vjp_ret
2644        else:
2645            y, f_vjp = vjp_ret
2646        if np.shape(y) != ():
2647            raise TypeError("grad output must be 0-dim scalar with shape ()")
2648        gL_xs = f_vjp(backend.ones(()))
2649        gL_xs = tuple(gL_xs[i] for i in argnums) if len(argnums) > 1 else gL_xs[argnums[0]]
2650        if return_value:
2651            return ((y, aux), gL_xs) if has_aux else (y, gL_xs)
2652        else:
2653            return (gL_xs, aux) if has_aux else gL_xs
2654
2655    return jit(gfn) if rejit else gfn
2656
2657
2658def value_and_grad(f, argnums=(0,), argnames="", has_aux=False):
2659    return grad(
2660        f,
2661        argnums=argnums,
2662        argnames=argnames,
2663        has_aux=has_aux,
2664        return_value=True,
2665    )
2666
2667
2668def jit_partial_run(trace, tracers, *, program):
2669    in_unknowns = [not t.pval.is_known for t in tracers]
2670    program1, program2, out_unknowns, num_res = partial_run_program(program, in_unknowns)
2671    known_tracers, unknown_tracers = partition_list(in_unknowns, tracers)
2672    known_vals = [t.pval.const for t in known_tracers]
2673    outs1_res = backend.jit_op(*known_vals, program=program)
2674    outs1, res = split_list(outs1_res, len(program1.outs) - num_res)
2675    res_tracers = [trace.instantiate_const(full_raise(trace, x)) for x in res]
2676    outs2 = [PartialRunTraceTensor(trace, PartialValue.unknown(v.symval), None) for v in program2.outs]
2677    draft = InstructionDraft(
2678        backend.jit_op,
2679        res_tracers + unknown_tracers,
2680        dict(program=program2),
2681        [v.symval for v in program2.outs],
2682        map(weakref.ref, outs2),
2683    )
2684    for t in outs2:
2685        t.draft = draft
2686    return merge_lists(out_unknowns, outs1, outs2)
2687
2688
2689class jit:
2690    def __init__(self, f, static_argnames=(), name=None, dynamic_axes=None):
2691        if isinstance(static_argnames, str):
2692            static_argnames = tuple(static_argnames.split(" "))
2693        assert type(static_argnames) is tuple and all(type(s) is str for s in static_argnames)
2694        self.f = f
2695        self.name = name if name is not None else self.f.__name__
2696        self.static_argnames = static_argnames
2697        self.dynamic_axes = dynamic_axes
2698
2699    @classmethod
2700    def with_options(cls, **kwargs):
2701        return partial(cls, **kwargs)
2702
2703    @classmethod
2704    def get_jit_name(cls, args, static_args, prefix="jit", short=False):
2705        name = f"{prefix}_"
2706        if short:
2707            static_args_tup = tuple(static_args.items())
2708            ids = repr(hash((prefix, args, static_args_tup)))[-4:]
2709            name = f"{prefix}_{ids}"
2710        else:
2711            for a in args:
2712                name += f"shape_{a.shape}_dtype_{a.dtype.name}_"
2713            for k, v in static_args.items():
2714                name += f"{k}_{v}_"
2715            name = name.replace("(", "L")
2716            name = name.replace(")", "R")
2717            name = name.replace(",", "C")
2718            name = name.replace(" ", "")
2719            name = name.replace(".", "D")
2720
2721        return name
2722
2723    def get_program(self, *args, **static_args):
2724        sig = inspect.signature(self.f)
2725        if all("*" not in repr(v) for v in sig.parameters.values()):
2726            args_strs = [k for k, v in sig.parameters.items() if k != "self" and k not in self.static_argnames]
2727            static_args_strs = [k for k, v in sig.parameters.items() if k != "self" and k in self.static_argnames]
2728
2729            if args:
2730                if len(args) > len(args_strs):
2731                    assert static_args_strs
2732                    args, rest = args[: len(args_strs)], args[len(args_strs) :]
2733                    new_static_args = {k: rest_arg for k, rest_arg in zip(static_args_strs, rest) if k not in static_args}
2734                    static_args = {**new_static_args, **static_args}
2735            else:
2736                args = tuple([static_args[k] if k in static_args else arg for k, arg in zip(args_strs, args)])
2737
2738        symvals_in = tree_map(lambda x: SymbolicTensor.like(get_symval(x)), args)
2739        static_args = tuple(static_args.items())
2740        if self.name is None:
2741            self.name = f"jit_{str(hash((self.f, symvals_in, static_args)))[-5:]}"
2742        program, consts, out_tree = make_program(self.f, *symvals_in, static_args=static_args, name=self.name)
2743        return program, consts, out_tree
2744
2745    def __call__(self, *args, **static_args):
2746        program, consts, out_tree = self.get_program(*args, **static_args)
2747        args, in_tree = tree_flatten(args)
2748        outs = bind(backend.jit_op, *consts, *args, program=program)
2749        return tree_unflatten(out_tree, outs)
2750
2751    def lower(self, *args, **static_args):
2752        program, consts, out_tree = self.get_program(*args, **static_args)
2753        args, in_tree = tree_flatten(args)
2754        hashed_program = Hashed(program)
2755        num_consts = program.num_consts
2756        consts, args = args[:num_consts], args[num_consts:]
2757        hashed_consts = tuple(map(Hashed, consts))
2758        jit_output = backend.jit_program(hashed_program, hashed_consts)
2759        return jit_output
2760
2761    def export(self, output_path, args, export_params=True, input_names=None, output_names=None, **kwargs):
2762        if isinstance(args, Tensor):
2763            args, static_args = (args,), dict()
2764        elif not isinstance(args[-1], dict):
2765            assert all(isinstance(a, Tensor) for a in args)
2766            static_args = dict()
2767        else:
2768            args, static_args = args[:-1], args[-1]
2769        assert isinstance(args, (tuple, list)) and isinstance(static_args, dict)
2770        jit_output = self.lower(*args, **static_args)
2771        backend.export(jit_output, output_path, export_params, input_names, output_names, **kwargs)