--- title: slope.core ---

# slope.core

View Source

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

    46class Timing(ContextDecorator):
    47    def __init__(self, prefix="", on_exit=None, enabled=True):
    48        self.prefix, self.on_exit, self.enabled = prefix, on_exit, enabled
    49
    50    def __enter__(self):
    51        self.st = time.perf_counter_ns()
    52
    53    def __exit__(self, *exc):
    54        self.et = time.perf_counter_ns() - self.st
    55        if self.enabled:
    56            print(f"{self.prefix}{self.et*1e-6:6.2f} ms" + (self.on_exit(self.et) if self.on_exit else ""))

A base class or mixin that enables context managers to work as decorators.

Timing(prefix='', on_exit=None, enabled=True) View Source

    47    def __init__(self, prefix="", on_exit=None, enabled=True):
    48        self.prefix, self.on_exit, self.enabled = prefix, on_exit, enabled
def colored(st, color: Optional[str], background=False): View Source

    59def colored(st, color: Optional[str], background=False):
    60    return (
    61        f"\u001b[{10*background+60*(color.upper() == color)+30+['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'].index(color.lower())}m{st}\u001b[0m"
    62        if color is not None
    63        else st
    64    )  # replace the termcolor library with one line  # noqa: E501
class Profiling(contextlib.ContextDecorator): View Source

    71class Profiling(ContextDecorator):
    72    def __init__(self, enabled=True, sort="cumtime", frac=0.2, fn=None, ts=1):
    73        self.enabled, self.sort, self.frac, self.fn, self.time_scale = enabled, sort, frac, fn, 1e3 / ts
    74
    75    def __enter__(self):
    76        self.pr = cProfile.Profile()
    77        if self.enabled:
    78            self.pr.enable()
    79
    80    def __exit__(self, *exc):
    81        if self.enabled:
    82            self.pr.disable()
    83            if self.fn:
    84                self.pr.dump_stats(self.fn)
    85            stats = pstats.Stats(self.pr).strip_dirs().sort_stats(self.sort)
    86            for fcn in stats.fcn_list[0 : int(len(stats.fcn_list) * self.frac)]:  # type: ignore[attr-defined]
    87                (_primitive_calls, num_calls, tottime, cumtime, callers) = stats.stats[fcn]  # type: ignore[attr-defined]
    88                scallers = sorted(callers.items(), key=lambda x: -x[1][2])
    89                print(
    90                    f"n:{num_calls:8d}  tm:{tottime*self.time_scale:7.2f}ms  tot:{cumtime*self.time_scale:7.2f}ms",
    91                    colored(_format_fcn(fcn), "yellow") + " " * (50 - len(_format_fcn(fcn))),
    92                    colored(f"<- {(scallers[0][1][2]/tottime)*100:3.0f}% {_format_fcn(scallers[0][0])}", "BLACK") if len(scallers) else "",
    93                )

A base class or mixin that enables context managers to work as decorators.

Profiling(enabled=True, sort='cumtime', frac=0.2, fn=None, ts=1) View Source

    72    def __init__(self, enabled=True, sort="cumtime", frac=0.2, fn=None, ts=1):
    73        self.enabled, self.sort, self.frac, self.fn, self.time_scale = enabled, sort, frac, fn, 1e3 / ts
def dblog(*msg, enable=True): View Source

    96def dblog(*msg, enable=True):
    97    if enable:
    98        print(*msg)
def unzip2(pairs) -> Tuple[List[Any], List[Any]]: View Source

    101def unzip2(pairs) -> Tuple[List[Any], List[Any]]:
    102    lst1, lst2 = [], []
    103    for i1, i2 in pairs:
    104        lst1 += [i1]
    105        lst2 += [i2]
    106    return lst1, lst2
def list_map(f: Callable, *xs: Iterable) -> List[Any]: View Source

    109def list_map(f: Callable, *xs: Iterable) -> List[Any]:
    110    return list(map(f, *xs))
def list_zip(*args: List[Any]) -> List[Any]: View Source

    113def list_zip(*args: List[Any]) -> List[Any]:
    114    fst, *rest = args = list_map(list, args)
    115    n = len(fst)
    116    for arg in rest:
    117        assert len(arg) == n
    118    return list(zip(*args))
def split_half(lst: List[Any]) -> Tuple[List[Any], List[Any]]: View Source

    121def split_half(lst: List[Any]) -> Tuple[List[Any], List[Any]]:
    122    assert not len(lst) % 2
    123    return split_list(lst, len(lst) // 2)
def merge_lists(which: List[bool], l1: List[Any], l2: List[Any]) -> List[Any]: View Source

    126def merge_lists(which: List[bool], l1: List[Any], l2: List[Any]) -> List[Any]:
    127    l1, l2 = iter(l1), iter(l2)
    128    out = [next(l2) if b else next(l1) for b in which]
    129    assert next(l1, None) is next(l2, None) is None
    130    return out
def split_list(lst: List[Any], n: int) -> Tuple[List[Any], List[Any]]: View Source

    148def split_list(lst: List[Any], n: int) -> Tuple[List[Any], List[Any]]:
    149    assert 0 <= n <= len(lst)
    150    return lst[:n], lst[n:]
def partition_list(bs: List[bool], l: List[Any]) -> Tuple[List[Any], List[Any]]: View Source

    153def partition_list(bs: List[bool], l: List[Any]) -> Tuple[List[Any], List[Any]]:
    154    assert len(bs) == len(l)
    155    lists = lst1, lst2 = [], []
    156    for b, x in zip(bs, l):
    157        lists[b].append(x)
    158    return lst1, lst2
def lru_cache_verbose( maxsize: int = 100, typed: bool = False, tb_start: int = -12, tb_end: int = -7): View Source

    161def lru_cache_verbose(
    162    maxsize: int = 100,
    163    typed: bool = False,
    164    tb_start: int = -12,
    165    tb_end: int = -7,
    166):
    167    def decorator(fn: Callable):
    168        @lru_cache(maxsize=maxsize, typed=typed)
    169        def wrapper(*args, **kwargs) -> Callable:
    170            return fn(*args, **kwargs)
    171
    172        def decorated_function(*args, **kwargs) -> Any:
    173            result = wrapper(*args, **kwargs)
    174            cache_info = wrapper.cache_info()
    175
    176            dblog(
    177                f"{fn.__name__}.{cache_info} {args.__hash__()}",
    178                enable=backend.LOG_LRU,
    179            )
    180            tb = "".join(traceback.format_list(traceback.extract_stack())[tb_start:tb_end]).replace("\n    ", ":\t") + "-" * 20 + "\n"
    181            dblog(f"{tb}", enable=backend.LOG_LRU)
    182
    183            return result
    184
    185        decorated_function.cache_info = wrapper.cache_info
    186        decorated_function.fn = fn
    187        return decorated_function
    188
    189    return decorator
def cuda_is_available(): View Source

    192def cuda_is_available():
    193    try:
    194        import subprocess
    195        import platform
    196
    197        cmd = f"nvidia-smi{'.exe' if platform.system == 'Windows' else ''}"
    198        result = subprocess.run([cmd], stdout=subprocess.PIPE)
    199        output = result.stdout.decode("utf-8")
    200        return True if "NVIDIA-SMI" in output else False
    201    except FileNotFoundError:
    202        return False
class Hashed: View Source

    205class Hashed:
    206    val: Any
    207
    208    def __init__(self, val):
    209        self.val = val
    210
    211    def __hash__(self) -> int:
    212        return hash((self.val,))
    213
    214    def __eq__(self, other):
    215        if isinstance(other, Hashed):
    216            if isinstance(self.val, Tensor) and isinstance(other.val, Tensor):
    217                # because Tensor.__eq__ already for Tensor.equal
    218                return id(self.val) == id(other.val)
    219            return self.val == other.val
    220        return False
    221
    222    def __repr__(self):
    223        return f"Hashed: {repr(self.val)}"
Hashed(val) View Source

    208    def __init__(self, val):
    209        self.val = val
val: Any
class DType(typing.NamedTuple): View Source

    231class DType(NamedTuple):
    232    priority: int
    233    itemsize: int
    234    name: str
    235    mlir: str
    236    numpy: type
    237
    238    @property
    239    def format_code(self):
    240        return f"slope.{self.name}"
    241
    242    def __repr__(self):
    243        return f"<DType: {self.name}>"

DType(priority, itemsize, name, mlir, numpy)

DType(priority: int, itemsize: int, name: str, mlir: str, numpy: type)

Create new instance of DType(priority, itemsize, name, mlir, numpy)

priority: int

Alias for field number 0

itemsize: int

Alias for field number 1

name: str

Alias for field number 2

mlir: str

Alias for field number 3

numpy: type

Alias for field number 4

format_code View Source

    238    @property
    239    def format_code(self):
    240        return f"slope.{self.name}"

##### Inherited Members

builtins.tuple
  index
  count

class dtypes: View Source

    246class dtypes:
    247    float32: Final[DType] = DType(4, 4, "float32", "f32", np.float32)
    248    uint8: Final[DType] = DType(0, 1, "uint8", "u8", np.uint8)
    249    int8: Final[DType] = DType(0, 1, "int8", "i8", np.int8)
    250    bool: Final[DType] = DType(0, 1, "bool", "i1", bool)
    251    int32: Final[DType] = DType(1, 4, "int32", "i32", np.int32)
    252    int64: Final[DType] = DType(2, 8, "int64", "i64", np.int64)
    253    uint64: Final[DType] = DType(2, 8, "uint64", "ui64", np.uint64)
    254    float16: Final[DType] = DType(0, 2, "float16", "f16", np.float16)
    255
    256    all_dtypes = (bool, float16, float32, int8, int32, int64, uint8, uint64)
    257    name_dtype_map = {k.name: k for k in all_dtypes}
    258    name_dtype_map_inv = {v: k for k, v in name_dtype_map.items()}
    259    mlir_dtype_map = {k.mlir: k for k in all_dtypes}
    260    mlir_dtype_map_inv = {v: k for k, v in mlir_dtype_map.items()}
    261
    262    @classmethod
    263    def is_int(cls, dtype):
    264        return dtype in (cls.uint8, cls.int8, cls.int32, cls.uint64, cls.int64)
    265
    266    @classmethod
    267    def is_float(cls, dtype):
    268        return dtype in (cls.float16, cls.float32)
float32: Final[DType] = <DType: float32>
uint8: Final[DType] = <DType: uint8>
int8: Final[DType] = <DType: int8>
bool: Final[DType] = <DType: bool>
int32: Final[DType] = <DType: int32>
int64: Final[DType] = <DType: int64>
uint64: Final[DType] = <DType: uint64>
float16: Final[DType] = <DType: float16>
all_dtypes = (<DType: bool>, <DType: float16>, <DType: float32>, <DType: int8>, <DType: int32>, <DType: int64>, <DType: uint8>, <DType: uint64>)
name_dtype_map = {'bool': <DType: bool>, 'float16': <DType: float16>, 'float32': <DType: float32>, 'int8': <DType: int8>, 'int32': <DType: int32>, 'int64': <DType: int64>, 'uint8': <DType: uint8>, 'uint64': <DType: uint64>}
name_dtype_map_inv = {<DType: bool>: 'bool', <DType: float16>: 'float16', <DType: float32>: 'float32', <DType: int8>: 'int8', <DType: int32>: 'int32', <DType: int64>: 'int64', <DType: uint8>: 'uint8', <DType: uint64>: 'uint64'}
mlir_dtype_map = {'i1': <DType: bool>, 'f16': <DType: float16>, 'f32': <DType: float32>, 'i8': <DType: int8>, 'i32': <DType: int32>, 'i64': <DType: int64>, 'u8': <DType: uint8>, 'ui64': <DType: uint64>}
mlir_dtype_map_inv = {<DType: bool>: 'i1', <DType: float16>: 'f16', <DType: float32>: 'f32', <DType: int8>: 'i8', <DType: int32>: 'i32', <DType: int64>: 'i64', <DType: uint8>: 'u8', <DType: uint64>: 'ui64'}
@classmethod
def is_int(cls, dtype): View Source

    262    @classmethod
    263    def is_int(cls, dtype):
    264        return dtype in (cls.uint8, cls.int8, cls.int32, cls.uint64, cls.int64)
@classmethod
def is_float(cls, dtype): View Source

    266    @classmethod
    267    def is_float(cls, dtype):
    268        return dtype in (cls.float16, cls.float32)
class Device(typing.NamedTuple): View Source

    271class Device(NamedTuple):
    272    name: str
    273    idx: int
    274
    275    @property
    276    def format_code(self):
    277        return f"'{self.name}:{self.idx}'"
    278
    279    def __repr__(self):
    280        return f"<Device: {self.format_code}>"

Device(name, idx)

Device(name: str, idx: int)

Create new instance of Device(name, idx)

name: str

Alias for field number 0

idx: int

Alias for field number 1

format_code View Source

    275    @property
    276    def format_code(self):
    277        return f"'{self.name}:{self.idx}'"

##### Inherited Members

builtins.tuple
  index
  count

class devices: View Source

    283class devices:
    284    cpu: Final[Device] = Device("cpu", 0)
    285    metal: Final[Device] = Device("metal", 0)
    286    cuda0: Final[Device] = Device("cuda", 0)
    287    # TODO: programmatically define this class attrs to support other setup
    288    cuda = cuda0
    289    all_devices = (cpu, metal, cuda0)
    290    name_idx_device_map = {f"{k.name}:{k.idx}": k for k in all_devices}
    291    name_idx_device_map_inv = {v: k for k, v in name_idx_device_map.items()}
cpu: Final[Device] = <Device: 'cpu:0'>
metal: Final[Device] = <Device: 'metal:0'>
cuda0: Final[Device] = <Device: 'cuda:0'>
cuda = <Device: 'cuda:0'>
all_devices = (<Device: 'cpu:0'>, <Device: 'metal:0'>, <Device: 'cuda:0'>)
name_idx_device_map = {'cpu:0': <Device: 'cpu:0'>, 'metal:0': <Device: 'metal:0'>, 'cuda:0': <Device: 'cuda:0'>}
name_idx_device_map_inv = {<Device: 'cpu:0'>: 'cpu:0', <Device: 'metal:0'>: 'metal:0', <Device: 'cuda:0'>: 'cuda:0'}
class TensorBuffer: View Source

    294class TensorBuffer:
    295    def __init__(self, val):
    296        self.val = val
TensorBuffer(val) View Source

    295    def __init__(self, val):
    296        self.val = val
val
class Tensor: View Source

    299class Tensor:
    300    def __init__(self, val: TensorBuffer):
    301        assert isinstance(val, TensorBuffer)
    302        self.buf = val
    303
    304    @property
    305    def symval(self):
    306        return SymbolicTensor.like(self)
    307
    308    @property
    309    def default_dtype(self):
    310        return backend.default_dtype
    311
    312    def is_int(self) -> bool:
    313        return self.dtype in (
    314            dtypes.int8,
    315            dtypes.uint8,
    316            dtypes.uint64,
    317            dtypes.int32,
    318            dtypes.int64,
    319        )
    320
    321    def is_float(self) -> bool:
    322        return self.dtype in (dtypes.float16, dtypes.float32)
    323
    324    def is_unsigned(self) -> bool:
    325        return self.dtype is dtypes.uint8
    326
    327    def to_bool(self):
    328        return self.cast(dtypes.bool)
    329
    330    def short(self):
    331        return self.cast(dtypes.int8)
    332
    333    def int(self):
    334        return self.cast(dtypes.int32)
    335
    336    def long(self):
    337        return self.cast(dtypes.int64)
    338
    339    def half(self):
    340        return self.cast(dtypes.float16)
    341
    342    def float(self):
    343        return self.cast(dtypes.float32)
    344
    345    def __getattr__(self, attr):
    346        if attr in vars(backend.operator_set).keys():
    347            op = getattr(backend.operator_set, attr)
    348            return partial(op, self)
    349        elif attr in vars(backend.procedure_set).keys():
    350            procedure = getattr(backend.procedure_set, attr)
    351            assert not isinstance(procedure, classmethod), f"use {attr} instead of self.{attr}"
    352            return partial(procedure, self)
    353        else:
    354            return self.__getattribute__(attr)
    355
    356    def __getitem__(self, idx):
    357        return self.getitem(idx)
    358
    359    def __setitem__(self, idx, item):
    360        raise NotImplementedError
    361
    362    def str_short(self):
    363        return f"<Tensor: shape={self.shape}, dtype={self.dtype}>"
    364
    365    __neg__ = lambda self: self.neg()
    366    __add__ = lambda self, other: self.add(other)
    367    __radd__ = lambda self, other: self.add(other)
    368    __sub__ = lambda self, other: self.sub(other)
    369    __rsub__ = lambda self, other: self.sub.func(other, self)
    370    __mul__ = lambda self, other: self.mul(other)
    371    __rmul__ = lambda self, other: self.mul(other)
    372    __div__ = lambda self, other: self.div(other)
    373    __rdiv__ = lambda self, other: self.div.func(other, self)
    374    __truediv__ = __div__
    375    __truerdiv__ = __rdiv__
    376    __pow__ = lambda self, other: self.pow(other)
    377    __rpow__ = lambda self, other: self.pow.func(other, self)
    378    __matmul__ = lambda self, other: self.matmul(other)
    379    __rmatmul__ = lambda self, other: self.matmul.func(other, self)
    380    __invert__ = lambda self: self.invert()
    381    __eq__ = lambda self, other: self.equal(other)
    382    __ne__ = lambda self, other: self.not_equal(other)
    383    __ge__ = lambda self, other: self.greater_equal(other)
    384    __le__ = lambda self, other: self.less_equal(other)
    385    __gt__ = lambda self, other: self.greater(other)
    386    __lt__ = lambda self, other: self.less(other)
    387
    388    def __hash__(self):
    389        return id(self.val)
    390
    391    val = property(lambda self: self.buf.val)
    392
    393    def size(self, i):
    394        return self.shape[i]
    395
    396    @property
    397    def dtype(self):
    398        return backend.dtype_of(self)
    399
    400    @property
    401    def device(self):
    402        return backend.device_of(self)
    403
    404    def numpy(self, memmap=False):
    405        return backend.numpy_of(self, memmap)
    406
    407    @property
    408    def shape(self):
    409        return backend.shape_of(self)
    410
    411    @property
    412    def ndim(self):
    413        return len(self.shape)
    414
    415    def numel(self):
    416        return math.prod(self.shape)
    417
    418    def element_size(self):
    419        return self.dtype.itemsize
    420
    421    def nbytes(self):
    422        return self.numel() * self.element_size()
    423
    424    def __repr__(self):
    425        return f"<Tensor: val=\n{self.numpy()}\nshape={self.shape}, dtype={self.dtype.name}, device={self.device.format_code}>"
Tensor(val: TensorBuffer) View Source

    300    def __init__(self, val: TensorBuffer):
    301        assert isinstance(val, TensorBuffer)
    302        self.buf = val
buf
symval View Source

    304    @property
    305    def symval(self):
    306        return SymbolicTensor.like(self)
default_dtype View Source

    308    @property
    309    def default_dtype(self):
    310        return backend.default_dtype
def is_int(self) -> bool: View Source

    312    def is_int(self) -> bool:
    313        return self.dtype in (
    314            dtypes.int8,
    315            dtypes.uint8,
    316            dtypes.uint64,
    317            dtypes.int32,
    318            dtypes.int64,
    319        )
def is_float(self) -> bool: View Source

    321    def is_float(self) -> bool:
    322        return self.dtype in (dtypes.float16, dtypes.float32)
def is_unsigned(self) -> bool: View Source

    324    def is_unsigned(self) -> bool:
    325        return self.dtype is dtypes.uint8
def to_bool(self): View Source

    327    def to_bool(self):
    328        return self.cast(dtypes.bool)
def short(self): View Source

    330    def short(self):
    331        return self.cast(dtypes.int8)
def int(self): View Source

    333    def int(self):
    334        return self.cast(dtypes.int32)
def long(self): View Source

    336    def long(self):
    337        return self.cast(dtypes.int64)
def half(self): View Source

    339    def half(self):
    340        return self.cast(dtypes.float16)
def float(self): View Source

    342    def float(self):
    343        return self.cast(dtypes.float32)
def str_short(self): View Source

    362    def str_short(self):
    363        return f"<Tensor: shape={self.shape}, dtype={self.dtype}>"
val View Source

    391    val = property(lambda self: self.buf.val)
def size(self, i): View Source

    393    def size(self, i):
    394        return self.shape[i]
dtype View Source

    396    @property
    397    def dtype(self):
    398        return backend.dtype_of(self)
device View Source

    400    @property
    401    def device(self):
    402        return backend.device_of(self)
def numpy(self, memmap=False): View Source

    404    def numpy(self, memmap=False):
    405        return backend.numpy_of(self, memmap)
shape View Source

    407    @property
    408    def shape(self):
    409        return backend.shape_of(self)
ndim View Source

    411    @property
    412    def ndim(self):
    413        return len(self.shape)
def numel(self): View Source

    415    def numel(self):
    416        return math.prod(self.shape)
def element_size(self): View Source

    418    def element_size(self):
    419        return self.dtype.itemsize
def nbytes(self): View Source

    421    def nbytes(self):
    422        return self.numel() * self.element_size()
class SymbolicTensor(Tensor): View Source

    428class SymbolicTensor(Tensor):
    429    def __init__(self, shape, dtype, device):
    430        assert isinstance(dtype, DType)
    431        self._shape = tuple(int(i) for i in shape)
    432        self._dtype = dtype
    433        self._device = device
    434
    435    @property
    436    def symval(self):
    437        return self
    438
    439    @property
    440    def val(self):
    441        raise RuntimeError(f"SymbolicTensor actually has no val, from {trace_stack[-1]=}, ")
    442
    443    @property
    444    def shape(self):
    445        return self._shape
    446
    447    @property
    448    def dtype(self):
    449        return self._dtype
    450
    451    @property
    452    def device(self):
    453        return self._device
    454
    455    def like(self, **but):
    456        shape = but.get("shape", self.shape)
    457        dtype = but.get("dtype", self.dtype)
    458        device = but.get("device", self.device)
    459        return SymbolicTensor(shape, dtype, device)
    460
    461    def str_short(self):
    462        return f'{str(self.dtype)}[{",".join(str(d) for d in self.shape)}]'
    463
    464    def __hash__(self):
    465        return hash((self.shape, self.dtype))
    466
    467    def __eq__(self, other):
    468        if type(self) != type(other):
    469            return False
    470        return tuple(self.shape) == tuple(other.shape) and self.dtype == other.dtype
    471
    472    def __repr__(self):
    473        return f"<SymbolicTensor: shape={self.shape}, dtype={self.dtype.name}, device={self.device}>"
SymbolicTensor(shape, dtype, device) View Source

    429    def __init__(self, shape, dtype, device):
    430        assert isinstance(dtype, DType)
    431        self._shape = tuple(int(i) for i in shape)
    432        self._dtype = dtype
    433        self._device = device
symval View Source

    435    @property
    436    def symval(self):
    437        return self
val View Source

    439    @property
    440    def val(self):
    441        raise RuntimeError(f"SymbolicTensor actually has no val, from {trace_stack[-1]=}, ")
shape View Source

    443    @property
    444    def shape(self):
    445        return self._shape
dtype View Source

    447    @property
    448    def dtype(self):
    449        return self._dtype
device View Source

    451    @property
    452    def device(self):
    453        return self._device
def like(self, **but): View Source

    455    def like(self, **but):
    456        shape = but.get("shape", self.shape)
    457        dtype = but.get("dtype", self.dtype)
    458        device = but.get("device", self.device)
    459        return SymbolicTensor(shape, dtype, device)
def str_short(self): View Source

    461    def str_short(self):
    462        return f'{str(self.dtype)}[{",".join(str(d) for d in self.shape)}]'

##### Inherited Members

Tensor
  buf
  default_dtype
  is_int
  is_float
  is_unsigned
  to_bool
  short
  int
  long
  half
  float
  size
  numpy
  ndim
  numel
  element_size
  nbytes

class Operator: View Source

    481class Operator:
    482    def __init__(self, name, variadic_inputs=False, nary_outputs=False):
    483        self.name = name
    484        self.variadic_inputs = variadic_inputs
    485        self.nary_outputs = nary_outputs
    486        if self.variadic_inputs:
    487            self.reorg_args = self.reorg_args_nary
    488
    489    def __hash__(self):
    490        return hash(self.name)
    491
    492    def __eq__(self, other):
    493        if not isinstance(other, Operator):
    494            return False
    495        return self.name == other.name
    496
    497    def args_fixer(self, *args, **params):
    498        return args, params
    499
    500    def __call__(self, *args, **params):
    501        args, params = self.reorg_args(args, params)
    502        args, params = self.args_fixer(*args, **params)
    503        ret = bind(self, *args, **params)
    504        if not self.nary_outputs:
    505            ret = ret[0]
    506        return ret
    507
    508    def __repr__(self) -> str:
    509        return f"<{self.name}>"
    510
    511    def typecheck(self, *args, **params):
    512        raise NotImplementedError
    513
    514    def jvp(self, *args, **params):
    515        raise NotImplementedError
    516
    517    def T(self, *args, **params):
    518        raise NotImplementedError
    519
    520    def vmap(self, *args, **params):
    521        raise NotImplementedError
    522
    523    def reorg_args(self, args, params):
    524        sig = inspect.signature(self.typecheck)
    525        args_strs = [k for k, v in sig.parameters.items() if v.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD and k != "self"]
    526        params_strs = [k for k, v in sig.parameters.items() if v.kind == inspect.Parameter.KEYWORD_ONLY and k != "self"]
    527
    528        if args:
    529            if len(args) > len(args_strs):
    530                args, rest = args[: len(args_strs)], args[len(args_strs) :]
    531                if params_strs:
    532                    new_params = {k: rest_arg for k, rest_arg in zip(params_strs, rest) if k not in params}
    533                    params = {**new_params, **params}
    534            else:
    535                args = tuple([params[k] if k in params else arg for k, arg in zip(args_strs, args)])
    536                assert len(args) == len(args_strs)
    537        return args, params
    538
    539    def reorg_args_nary(self, args, params):
    540        return args, params
    541
    542    def partial_run(self, trace, tracers, **params):
    543        tracers_in = [trace.instantiate_const(t) for t in tracers]
    544        symvals_in = [t.symval for t in tracers_in]
    545        symvals_out = self.typecheck(*symvals_in, **params)
    546        tracers_out = [PartialRunTraceTensor(trace, make_unknown_pval(symval), None) for symval in symvals_out]
    547        instruction = InstructionDraft(
    548            self,
    549            tracers_in,
    550            params,
    551            symvals_out,
    552            list_map(weakref.ref, tracers_out),
    553        )
    554        for t in tracers_out:
    555            t.draft = instruction
    556        return tracers_out
    557
    558    def partial_run_instruction(self, unks_in, instruction):
    559        if any(unks_in):
    560            instruction1 = None
    561            instruction2 = Instruction(
    562                instruction.op,
    563                instruction.inputs,
    564                instruction.params,
    565                instruction.out_binders,
    566            )
    567            unks_out = [True for i in instruction.out_binders]
    568            res = [v for unk, v in zip(unks_in, instruction.inputs) if ((not unk) and type(v) is Var)]
    569        else:
    570            instruction1 = instruction
    571            instruction2 = None
    572            unks_out = [False for i in instruction.out_binders]
    573            res = None
    574
    575        return instruction1, instruction2, unks_out, res
Operator(name, variadic_inputs=False, nary_outputs=False) View Source

    482    def __init__(self, name, variadic_inputs=False, nary_outputs=False):
    483        self.name = name
    484        self.variadic_inputs = variadic_inputs
    485        self.nary_outputs = nary_outputs
    486        if self.variadic_inputs:
    487            self.reorg_args = self.reorg_args_nary
name
variadic_inputs
nary_outputs
def args_fixer(self, *args, **params): View Source

    497    def args_fixer(self, *args, **params):
    498        return args, params
def typecheck(self, *args, **params): View Source

    511    def typecheck(self, *args, **params):
    512        raise NotImplementedError
def jvp(self, *args, **params): View Source

    514    def jvp(self, *args, **params):
    515        raise NotImplementedError
def T(self, *args, **params): View Source

    517    def T(self, *args, **params):
    518        raise NotImplementedError
def vmap(self, *args, **params): View Source

    520    def vmap(self, *args, **params):
    521        raise NotImplementedError
def reorg_args(self, args, params): View Source

    523    def reorg_args(self, args, params):
    524        sig = inspect.signature(self.typecheck)
    525        args_strs = [k for k, v in sig.parameters.items() if v.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD and k != "self"]
    526        params_strs = [k for k, v in sig.parameters.items() if v.kind == inspect.Parameter.KEYWORD_ONLY and k != "self"]
    527
    528        if args:
    529            if len(args) > len(args_strs):
    530                args, rest = args[: len(args_strs)], args[len(args_strs) :]
    531                if params_strs:
    532                    new_params = {k: rest_arg for k, rest_arg in zip(params_strs, rest) if k not in params}
    533                    params = {**new_params, **params}
    534            else:
    535                args = tuple([params[k] if k in params else arg for k, arg in zip(args_strs, args)])
    536                assert len(args) == len(args_strs)
    537        return args, params
def reorg_args_nary(self, args, params): View Source

    539    def reorg_args_nary(self, args, params):
    540        return args, params
def partial_run(self, trace, tracers, **params): View Source

    542    def partial_run(self, trace, tracers, **params):
    543        tracers_in = [trace.instantiate_const(t) for t in tracers]
    544        symvals_in = [t.symval for t in tracers_in]
    545        symvals_out = self.typecheck(*symvals_in, **params)
    546        tracers_out = [PartialRunTraceTensor(trace, make_unknown_pval(symval), None) for symval in symvals_out]
    547        instruction = InstructionDraft(
    548            self,
    549            tracers_in,
    550            params,
    551            symvals_out,
    552            list_map(weakref.ref, tracers_out),
    553        )
    554        for t in tracers_out:
    555            t.draft = instruction
    556        return tracers_out
def partial_run_instruction(self, unks_in, instruction): View Source

    558    def partial_run_instruction(self, unks_in, instruction):
    559        if any(unks_in):
    560            instruction1 = None
    561            instruction2 = Instruction(
    562                instruction.op,
    563                instruction.inputs,
    564                instruction.params,
    565                instruction.out_binders,
    566            )
    567            unks_out = [True for i in instruction.out_binders]
    568            res = [v for unk, v in zip(unks_in, instruction.inputs) if ((not unk) and type(v) is Var)]
    569        else:
    570            instruction1 = instruction
    571            instruction2 = None
    572            unks_out = [False for i in instruction.out_binders]
    573            res = None
    574
    575        return instruction1, instruction2, unks_out, res
class MetaOperator(Operator): View Source

    578class MetaOperator(Operator):
    579    def meta_impl(self, *args, **kwargs):
    580        raise NotImplementedError
def meta_impl(self, *args, **kwargs): View Source

    579    def meta_impl(self, *args, **kwargs):
    580        raise NotImplementedError

##### Inherited Members

Operator
  Operator
  name
  variadic_inputs
  nary_outputs
  args_fixer
  typecheck
  jvp
  T
  vmap
  reorg_args
  reorg_args_nary
  partial_run
  partial_run_instruction

class UnaryOperator(Operator): View Source

    583class UnaryOperator(Operator):
    584    def vmap(self, x, *, dim_size, vals_in, dims_in, **params):
    585        (x,), (x_bdim,) = vals_in, dims_in
    586        return [self(x, **params)], [x_bdim]
    587
    588    def typecheck(self, x, **params):
    589        return [x.like()]
    590
    591    def jvp(self, primals, tangents, **params):
    592        (x,), (x_dot,) = primals, tangents
    593        return [self(x, **params)], [self(x_dot, **params)]
def vmap(self, x, *, dim_size, vals_in, dims_in, **params): View Source

    584    def vmap(self, x, *, dim_size, vals_in, dims_in, **params):
    585        (x,), (x_bdim,) = vals_in, dims_in
    586        return [self(x, **params)], [x_bdim]
def typecheck(self, x, **params): View Source

    588    def typecheck(self, x, **params):
    589        return [x.like()]
def jvp(self, primals, tangents, **params): View Source

    591    def jvp(self, primals, tangents, **params):
    592        (x,), (x_dot,) = primals, tangents
    593        return [self(x, **params)], [self(x_dot, **params)]

##### Inherited Members

Operator
  Operator
  name
  variadic_inputs
  nary_outputs
  args_fixer
  T
  reorg_args
  reorg_args_nary
  partial_run
  partial_run_instruction

class BinaryOperator(Operator): View Source

    596class BinaryOperator(Operator):
    597    boolean_output = False
    598
    599    def args_fixer(self, x, w, **params):
    600        if isinstance(x, UndefinedPrimal) or type(w) is UndefinedPrimal:
    601            assert x.shape == w.shape
    602            return (x, w), params
    603
    604        if type(x) in TraceTensor.PYTHON_TYPES:
    605            x = backend.full(shape=(), fill_value=x, dtype=w.dtype)
    606        elif type(w) in TraceTensor.PYTHON_TYPES:
    607            w = backend.full(shape=(), fill_value=w, dtype=x.dtype)
    608
    609        shape_delta = x.ndim - w.ndim
    610        if shape_delta > 0:
    611            w = w.reshape((1,) * shape_delta + w.shape)
    612        elif shape_delta < 0:
    613            x = x.reshape((1,) * -shape_delta + x.shape)
    614
    615        shape_ret = tuple([max(x, w) for x, w in zip(x.shape, w.shape)])
    616        if x.shape != shape_ret:
    617            x = x.expand(shape_ret)
    618        if w.shape != shape_ret:
    619            w = w.expand(shape_ret)
    620
    621        if type(x) is Tensor and isinstance(w, TraceTensor):
    622            x = w._trace.pure(x)
    623        elif type(w) is Tensor and isinstance(x, TraceTensor):
    624            w = x._trace.pure(w)
    625        # TODO: https://jax.readthedocs.io/en/latest/type_promotion.html
    626        if x.dtype != w.dtype:
    627            # {int, bool} -> float
    628            if dtypes.is_float(x.dtype) ^ dtypes.is_float(w.dtype):
    629                if dtypes.is_float(w.dtype):
    630                    x = x.cast(w.dtype)
    631                elif dtypes.is_float(x.dtype):
    632                    w = w.cast(x.dtype)
    633            # bool -> int
    634            elif dtypes.is_int(x.dtype) ^ dtypes.is_int(w.dtype):
    635                if dtypes.is_int(w.dtype):
    636                    x = x.cast(w.dtype)
    637                elif dtypes.is_int(x.dtype):
    638                    w = w.cast(x.dtype)
    639            else:  # TODO: fine-grained type promotions
    640                raise NotImplementedError("No other type promotion rules")
    641
    642        return (x, w), params
    643
    644    def vmap(self, dim_size, vals_in, dims_in, **params):
    645        (x, w), (x_bdim, w_bdim) = vals_in, dims_in
    646        if x_bdim != w_bdim:
    647            if x_bdim is None:
    648                x = VMapTrace.move_vmap_dim(x, dim_size, x_bdim, w_bdim)
    649                x_bdim = w_bdim
    650            else:
    651                w = VMapTrace.move_vmap_dim(w, dim_size, w_bdim, x_bdim)
    652        return [self(x, w, **params)], [x_bdim]
    653
    654    def typecheck(self, x: SymbolicTensor, y: SymbolicTensor, **params) -> List[SymbolicTensor]:
    655        if not type(x) in (Tensor, SymbolicTensor) or not type(x) in (
    656            Tensor,
    657            SymbolicTensor,
    658        ):
    659            raise TypeError
    660        symx = SymbolicTensor.like(x, dtype=dtypes.bool if self.boolean_output else x.dtype)
    661        symy = SymbolicTensor.like(y, dtype=dtypes.bool if self.boolean_output else y.dtype)
    662        if x.dtype != y.dtype:
    663            raise TypeError(f"x.dtype ({x.dtype}) != y.dtype ({y.dtype})")
    664        if symx == symy:
    665            return [symx]
    666        shape_delta = len(symx.shape) - len(symy.shape)
    667        if shape_delta > 0:
    668            symy = symy.like(shape=(1,) * shape_delta + symy.shape)
    669        elif shape_delta < 0:
    670            x = x.reshape((1,) * -shape_delta + symx.shape)
    671            symx = symx.like(shape=(1,) * -shape_delta + symx.shape)
    672        if symx == symy:
    673            return [symx]
    674        else:
    675            shape_ret = tuple([max(x, w) for x, w in zip(symx.shape, symy.shape)])
    676            if symx.shape != shape_ret:
    677                symx = symx.like(shape=shape_ret)
    678            if symy.shape != shape_ret:
    679                symy = symx.like(shape=shape_ret)
    680            if symx != symy:
    681                raise TypeError(f"symx ({symx}) != symy ({symy})")
    682            return [symx]
    683
    684    def jvp(self, primals, tangents, **params):
    685        (x, w), (x_dot, w_dot) = primals, tangents
    686        return [self(x, w, **params)], [self(x_dot, w_dot, **params)]
    687
    688    def T(self, cotangents, x, w):
    689        (gL_y,) = cotangents
    690        if self.boolean_output:
    691            gL_y = gL_y.cast(x.dtype)
    692        if isinstance(x, UndefinedPrimal):
    693            return [gL_y, NullCotangent]
    694        elif type(w) is UndefinedPrimal:
    695            return [NullCotangent, gL_y]
boolean_output = False
def args_fixer(self, x, w, **params): View Source

    599    def args_fixer(self, x, w, **params):
    600        if isinstance(x, UndefinedPrimal) or type(w) is UndefinedPrimal:
    601            assert x.shape == w.shape
    602            return (x, w), params
    603
    604        if type(x) in TraceTensor.PYTHON_TYPES:
    605            x = backend.full(shape=(), fill_value=x, dtype=w.dtype)
    606        elif type(w) in TraceTensor.PYTHON_TYPES:
    607            w = backend.full(shape=(), fill_value=w, dtype=x.dtype)
    608
    609        shape_delta = x.ndim - w.ndim
    610        if shape_delta > 0:
    611            w = w.reshape((1,) * shape_delta + w.shape)
    612        elif shape_delta < 0:
    613            x = x.reshape((1,) * -shape_delta + x.shape)
    614
    615        shape_ret = tuple([max(x, w) for x, w in zip(x.shape, w.shape)])
    616        if x.shape != shape_ret:
    617            x = x.expand(shape_ret)
    618        if w.shape != shape_ret:
    619            w = w.expand(shape_ret)
    620
    621        if type(x) is Tensor and isinstance(w, TraceTensor):
    622            x = w._trace.pure(x)
    623        elif type(w) is Tensor and isinstance(x, TraceTensor):
    624            w = x._trace.pure(w)
    625        # TODO: https://jax.readthedocs.io/en/latest/type_promotion.html
    626        if x.dtype != w.dtype:
    627            # {int, bool} -> float
    628            if dtypes.is_float(x.dtype) ^ dtypes.is_float(w.dtype):
    629                if dtypes.is_float(w.dtype):
    630                    x = x.cast(w.dtype)
    631                elif dtypes.is_float(x.dtype):
    632                    w = w.cast(x.dtype)
    633            # bool -> int
    634            elif dtypes.is_int(x.dtype) ^ dtypes.is_int(w.dtype):
    635                if dtypes.is_int(w.dtype):
    636                    x = x.cast(w.dtype)
    637                elif dtypes.is_int(x.dtype):
    638                    w = w.cast(x.dtype)
    639            else:  # TODO: fine-grained type promotions
    640                raise NotImplementedError("No other type promotion rules")
    641
    642        return (x, w), params
def vmap(self, dim_size, vals_in, dims_in, **params): View Source

    644    def vmap(self, dim_size, vals_in, dims_in, **params):
    645        (x, w), (x_bdim, w_bdim) = vals_in, dims_in
    646        if x_bdim != w_bdim:
    647            if x_bdim is None:
    648                x = VMapTrace.move_vmap_dim(x, dim_size, x_bdim, w_bdim)
    649                x_bdim = w_bdim
    650            else:
    651                w = VMapTrace.move_vmap_dim(w, dim_size, w_bdim, x_bdim)
    652        return [self(x, w, **params)], [x_bdim]
def typecheck( self, x: SymbolicTensor, y: SymbolicTensor, **params) -> List[SymbolicTensor]: View Source

    654    def typecheck(self, x: SymbolicTensor, y: SymbolicTensor, **params) -> List[SymbolicTensor]:
    655        if not type(x) in (Tensor, SymbolicTensor) or not type(x) in (
    656            Tensor,
    657            SymbolicTensor,
    658        ):
    659            raise TypeError
    660        symx = SymbolicTensor.like(x, dtype=dtypes.bool if self.boolean_output else x.dtype)
    661        symy = SymbolicTensor.like(y, dtype=dtypes.bool if self.boolean_output else y.dtype)
    662        if x.dtype != y.dtype:
    663            raise TypeError(f"x.dtype ({x.dtype}) != y.dtype ({y.dtype})")
    664        if symx == symy:
    665            return [symx]
    666        shape_delta = len(symx.shape) - len(symy.shape)
    667        if shape_delta > 0:
    668            symy = symy.like(shape=(1,) * shape_delta + symy.shape)
    669        elif shape_delta < 0:
    670            x = x.reshape((1,) * -shape_delta + symx.shape)
    671            symx = symx.like(shape=(1,) * -shape_delta + symx.shape)
    672        if symx == symy:
    673            return [symx]
    674        else:
    675            shape_ret = tuple([max(x, w) for x, w in zip(symx.shape, symy.shape)])
    676            if symx.shape != shape_ret:
    677                symx = symx.like(shape=shape_ret)
    678            if symy.shape != shape_ret:
    679                symy = symx.like(shape=shape_ret)
    680            if symx != symy:
    681                raise TypeError(f"symx ({symx}) != symy ({symy})")
    682            return [symx]
def jvp(self, primals, tangents, **params): View Source

    684    def jvp(self, primals, tangents, **params):
    685        (x, w), (x_dot, w_dot) = primals, tangents
    686        return [self(x, w, **params)], [self(x_dot, w_dot, **params)]
def T(self, cotangents, x, w): View Source

    688    def T(self, cotangents, x, w):
    689        (gL_y,) = cotangents
    690        if self.boolean_output:
    691            gL_y = gL_y.cast(x.dtype)
    692        if isinstance(x, UndefinedPrimal):
    693            return [gL_y, NullCotangent]
    694        elif type(w) is UndefinedPrimal:
    695            return [NullCotangent, gL_y]

##### Inherited Members

Operator
  Operator
  name
  variadic_inputs
  nary_outputs
  reorg_args
  reorg_args_nary
  partial_run
  partial_run_instruction

class ReduceOperator(Operator): View Source

    698class ReduceOperator(Operator):
    699    def args_fixer(self, x, *, dim=None, keepdim=False):
    700        if dim is None:
    701            dim = tuple(range(x.ndim))
    702        elif isinstance(dim, int):
    703            dim = (dim,)
    704        dim = tuple(a if a >= 0 else a + len(x.shape) for a in dim)
    705        return (x,), dict(dim=dim, keepdim=keepdim)
    706
    707    def vmap(self, dim_size, vals_in, dims_in, *, dim, keepdim):
    708        (x,), (x_bdim,) = vals_in, dims_in
    709        dim = tuple(a + (x_bdim <= a) for a in dim)
    710        out_bdim = x_bdim - sum(a < x_bdim for a in dim)
    711        return [self(x, dim=dim, keepdim=keepdim)], [out_bdim]
    712
    713    def typecheck(self, x: SymbolicTensor, *, dim=None, keepdim=False) -> List[SymbolicTensor]:
    714        dim = [a + len(x.shape) if a < 0 else a for a in dim]
    715        dim_ = set(dim)
    716        if keepdim:
    717            new_shape = [d if i not in dim_ else 1 for i, d in enumerate(x.shape)]
    718        else:
    719            new_shape = [d for i, d in enumerate(x.shape) if i not in dim_]
    720        return [x.like(shape=tuple(new_shape))]
def args_fixer(self, x, *, dim=None, keepdim=False): View Source

    699    def args_fixer(self, x, *, dim=None, keepdim=False):
    700        if dim is None:
    701            dim = tuple(range(x.ndim))
    702        elif isinstance(dim, int):
    703            dim = (dim,)
    704        dim = tuple(a if a >= 0 else a + len(x.shape) for a in dim)
    705        return (x,), dict(dim=dim, keepdim=keepdim)
def vmap(self, dim_size, vals_in, dims_in, *, dim, keepdim): View Source

    707    def vmap(self, dim_size, vals_in, dims_in, *, dim, keepdim):
    708        (x,), (x_bdim,) = vals_in, dims_in
    709        dim = tuple(a + (x_bdim <= a) for a in dim)
    710        out_bdim = x_bdim - sum(a < x_bdim for a in dim)
    711        return [self(x, dim=dim, keepdim=keepdim)], [out_bdim]
def typecheck( self, x: SymbolicTensor, *, dim=None, keepdim=False) -> List[SymbolicTensor]: View Source

    713    def typecheck(self, x: SymbolicTensor, *, dim=None, keepdim=False) -> List[SymbolicTensor]:
    714        dim = [a + len(x.shape) if a < 0 else a for a in dim]
    715        dim_ = set(dim)
    716        if keepdim:
    717            new_shape = [d if i not in dim_ else 1 for i, d in enumerate(x.shape)]
    718        else:
    719            new_shape = [d for i, d in enumerate(x.shape) if i not in dim_]
    720        return [x.like(shape=tuple(new_shape))]

##### Inherited Members

Operator
  Operator
  name
  variadic_inputs
  nary_outputs
  jvp
  T
  reorg_args
  reorg_args_nary
  partial_run
  partial_run_instruction

class InitOperator(Operator): View Source

    723class InitOperator(Operator):
    724    def vmap(self, dim_size, vals_in, dims_in, **params):
    725        (x_bdim,) = dims_in
    726        y = self(**params)
    727        y = y.unsqueeze(x_bdim)
    728        return [y], [x_bdim]
    729
    730    def jvp(self, primals, tangents, **params):
    731        y = self(**params)
    732        y_dot = NullCotangent(y.symval)
    733        # y_dot = y.ones_like()
    734        return [y], [y_dot]
    735
    736    def T(self, cotangents, **params):
    737        return [NullCotangent(cotangents[0])]
def vmap(self, dim_size, vals_in, dims_in, **params): View Source

    724    def vmap(self, dim_size, vals_in, dims_in, **params):
    725        (x_bdim,) = dims_in
    726        y = self(**params)
    727        y = y.unsqueeze(x_bdim)
    728        return [y], [x_bdim]
def jvp(self, primals, tangents, **params): View Source

    730    def jvp(self, primals, tangents, **params):
    731        y = self(**params)
    732        y_dot = NullCotangent(y.symval)
    733        # y_dot = y.ones_like()
    734        return [y], [y_dot]
def T(self, cotangents, **params): View Source

    736    def T(self, cotangents, **params):
    737        return [NullCotangent(cotangents[0])]

##### Inherited Members

Operator
  Operator
  name
  variadic_inputs
  nary_outputs
  args_fixer
  typecheck
  reorg_args
  reorg_args_nary
  partial_run
  partial_run_instruction

class ShapeOperator(Operator): View Source

    740class ShapeOperator(Operator):
    741    pass

##### Inherited Members

Operator
  Operator
  name
  variadic_inputs
  nary_outputs
  args_fixer
  typecheck
  jvp
  T
  vmap
  reorg_args
  reorg_args_nary
  partial_run
  partial_run_instruction

class GeneralReduceOperator(Operator): View Source

    744class GeneralReduceOperator(Operator):
    745    pass

##### Inherited Members

Operator
  Operator
  name
  variadic_inputs
  nary_outputs
  args_fixer
  typecheck
  jvp
  T
  vmap
  reorg_args
  reorg_args_nary
  partial_run
  partial_run_instruction

class OperatorSet: View Source

    748class OperatorSet:
    749    def __init__(self):
    750        self.register("jit_op")(JitOp)
    751
    752    def register(self, name, variadic_inputs=False, nary_outputs=False, aliases=()):
    753        def wrap(op_cls):
    754            assert name not in vars(self)
    755            op = op_cls(name, variadic_inputs, nary_outputs)
    756            setattr(self, name, op)
    757            for a in aliases:
    758                setattr(self, a, op)
    759            return op_cls
    760
    761        return wrap
def register(self, name, variadic_inputs=False, nary_outputs=False, aliases=()): View Source

    752    def register(self, name, variadic_inputs=False, nary_outputs=False, aliases=()):
    753        def wrap(op_cls):
    754            assert name not in vars(self)
    755            op = op_cls(name, variadic_inputs, nary_outputs)
    756            setattr(self, name, op)
    757            for a in aliases:
    758                setattr(self, a, op)
    759            return op_cls
    760
    761        return wrap
class ProcedureSet: View Source

    764class ProcedureSet:
    765    def register(self, aliases=()):
    766        def wrap(f):
    767            assert f.__name__ not in vars(self)
    768            setattr(self, f.__name__, f)
    769            for a in aliases:
    770                setattr(self, a, f)
    771            return f
    772
    773        return wrap
def register(self, aliases=()): View Source

    765    def register(self, aliases=()):
    766        def wrap(f):
    767            assert f.__name__ not in vars(self)
    768            setattr(self, f.__name__, f)
    769            for a in aliases:
    770                setattr(self, a, f)
    771            return f
    772
    773        return wrap
class CodegenOutput(typing.NamedTuple): View Source

    776class CodegenOutput(NamedTuple):
    777    code_lines: List[str]
    778    fn_defs: Dict[str, List[str]]
    779    in_binders: List["ProgramEnvVar"]
    780    outs: List["ProgramEnvVar"]

CodegenOutput(code_lines, fn_defs, in_binders, outs)

CodegenOutput( code_lines: List[str], fn_defs: Dict[str, List[str]], in_binders: List[ForwardRef('ProgramEnvVar')], outs: List[ForwardRef('ProgramEnvVar')])

Create new instance of CodegenOutput(code_lines, fn_defs, in_binders, outs)

code_lines: List[str]

Alias for field number 0

fn_defs: Dict[str, List[str]]

Alias for field number 1

in_binders: List[ProgramEnvVar]

Alias for field number 2

outs: List[ProgramEnvVar]

Alias for field number 3

##### Inherited Members

builtins.tuple
  index
  count

class Backend: View Source

    783class Backend:
    784    LOG_LRU = int(os.environ.get("LOG_LRU", 0))
    785    LOG_JIT = int(os.environ.get("LOG_JIT", 0))
    786    LOG_TREE = int(os.environ.get("LOG_TREE", 0))
    787    LOG_BACKEND = int(os.environ.get("LOG_BACKEND", 0))
    788    LOG_PROGRAM = int(os.environ.get("LOG_PROGRAM", 0))
    789    LOG_INIT = int(os.environ.get("LOG_INIT", 1))
    790    device_var = os.environ.get("DEFAULT_DEVICE", "cpu:0")
    791    if device_var[-2] != ":":
    792        device_var += ":0"
    793    DEFAULT_DEVICE = devices.name_idx_device_map[device_var]
    794    DEFAULT_DTYPE = dtypes.name_dtype_map[os.environ.get("DEFAULT_DTYPE", "float32")]
    795    dtype_for_indices: DType = None  # need to override
    796
    797    def __init__(
    798        self,
    799        operator_set: OperatorSet,
    800        procedure_set: ProcedureSet,
    801    ):
    802        self.operator_set = operator_set
    803        self.procedure_set = procedure_set
    804        self.node_types = dict()
    805        self.impls = dict()
    806        self.register_node(tuple, lambda t: (None, t), lambda _, xs: tuple(xs), "tuple")
    807        self.register_node(list, lambda l: (None, l), lambda _, xs: list(xs), "list")
    808        self.register_node(
    809            dict,
    810            lambda d: list_map(tuple, unzip2(sorted(d.items()))),
    811            lambda keys, vals: dict(list_zip(keys, vals)),
    812            "dict",
    813        )
    814        self.register_node(
    815            UndefinedPrimal,
    816            lambda u: (u.symval, ()),
    817            lambda symval, _: UndefinedPrimal(symval),
    818            "UndefinedPrimal",
    819        )
    820
    821    def set_impl(self, op: Union[types.LambdaType, types.FunctionType]):
    822        def set_impl_(fn):
    823            self.impls[op] = types.MethodType(fn, self)
    824
    825        return set_impl_
    826
    827    def register_node(self, ty: Type, to_iter: Callable, from_iter: Callable, name=None) -> None:
    828        if name is None:
    829            name = str(ty)
    830        self.node_types[ty] = NodeType(name, to_iter, from_iter)
    831
    832    def __getattr__(self, attr):
    833        try:
    834            dblog(
    835                f"Looking {self}.{attr} in operator_set",
    836                enable=backend.LOG_BACKEND,
    837            )
    838            return getattr(self.operator_set, attr)
    839        except:
    840            pass
    841        try:
    842            dblog(
    843                f"Looking {self}.{attr} in procedure_set",
    844                enable=backend.LOG_BACKEND,
    845            )
    846            return getattr(self.procedure_set, attr)
    847        except:
    848            pass
    849        dblog(
    850            f"Fallback to default {self} getattribute",
    851            enable=backend.LOG_BACKEND,
    852        )
    853        super().__getattribute__(attr)
    854
    855    def tensor(
    856        self,
    857        val: Union[list, tuple, np.ndarray, "TensorBuffer"] = None,
    858        dtype: Optional[Any] = None,
    859        device=None,
    860    ):
    861        if isinstance(val, TensorBuffer):
    862            return Tensor(val)
    863        elif isinstance(val, Tensor):
    864            return val
    865        if type(val) is bytes:
    866            val = np.frombuffer(val, dtype=dtype)
    867        return self.from_numpy(val, dtype, device)
    868
    869    def symbolic_tensor(
    870        self,
    871        shape: Union[list, tuple, np.ndarray, "TensorBuffer"] = None,
    872        dtype: Optional[Any] = None,
    873        device=None,
    874    ):
    875        dtype = dtype or self.DEFAULT_DTYPE
    876        device = device or self.DEFAULT_DEVICE
    877        return SymbolicTensor(shape, dtype, device)
    878
    879    def seed(self, seed):
    880        raise NotImplementedError
    881
    882    @property
    883    def default_dtype_value(self):
    884        return self.dtype_map[backend.DEFAULT_DTYPE]
    885
    886    def set_method(self, method):
    887        setattr(self, method.__name__, types.MethodType(method, self))
    888
    889    def from_numpy(self, val, device):
    890        raise NotImplementedError
    891
    892    def numpy_of(self, tensor):
    893        raise NotImplementedError
    894
    895    def device_of(self, tensor):
    896        raise NotImplementedError
    897
    898    def shape_of(self, tensor):
    899        raise NotImplementedError
    900
    901    def dtype_of(self, tensor):
    902        raise NotImplementedError
    903
    904    @lru_cache_verbose()
    905    def jit_program(
    906        self,
    907        hashed_program: Hashed,
    908        hashed_consts: Tuple[Hashed, ...],
    909    ):
    910        program: Program = hashed_program.val
    911        typecheck_program(program)
    912        consts = [x.val for x in hashed_consts]
    913        in_symvals = [v.symval for v in program.in_binders[len(consts) :]]
    914        codegen_output: CodegenOutput = self.codegen(program, consts + in_symvals, fn_name="main")
    915        fn, code = self.compile(codegen_output)
    916        jit_output = JitOutput(program, codegen_output, fn, code, consts)
    917        return jit_output
    918
    919    def codegen(self, program: "Program", args: Tuple, in_symvals: Tuple, name: str):
    920        "Returns compiler IR from the Program"
    921        raise NotImplementedError
    922
    923    def compile(self, program: "Program", args: Tuple, in_symvals: Tuple, name: str):
    924        "Compiles compiler IR to a Python callable function"
    925        raise NotImplementedError
    926
    927    def export(self, jit_output, *args, **params):
    928        raise NotImplementedError
    929
    930    def load(self, path, single_key="_tensor"):
    931        with open(path, mode="rb") as f:
    932            with mmap.mmap(f.fileno(), length=0, access=mmap.ACCESS_READ) as m:
    933                json_len = np.int64(m[0])
    934                start = 8 + json_len
    935                metadata = json.loads(m[8:start])
    936                ret = {}
    937                for k, v in metadata.items():
    938                    if k != "__metadata__":
    939                        dtype = Tensor.mlir_dtype_map[(v["dtype"])]
    940                        data_start = start + v["data_offsets"][0]
    941                        data_end = start + v["data_offsets"][1]
    942                        t_np = np.frombuffer(m[data_start:data_end], dtype=dtype.numpy())
    943                        t = backend.tensor(t_np, dtype=dtype)
    944                        t = t.reshape(tuple(v["shape"]))
    945                        ret[k] = t
    946                if len(ret) == 1 and single_key in ret.keys():
    947                    return ret[single_key]
    948                return ret
    949
    950    def save(self, tensors: Dict[str, Tensor], path: str, single_key="_tensor"):
    951        if isinstance(tensors, Tensor):
    952            tensors = {single_key: tensors}
    953        else:
    954            assert all((isinstance(k, str) and isinstance(v, Tensor)) for k, v in tensors.items())
    955
    956        metadata, offset = {}, 0
    957        for k, v in tensors.items():
    958            metadata[k] = {
    959                "dtype": v.dtype.mlir,
    960                "shape": list(v.shape),
    961                "data_offsets": [offset, offset + v.nbytes()],
    962            }
    963            offset += v.nbytes()
    964        j = json.dumps(metadata, separators=(",", ":"))
    965        Path(path).unlink(missing_ok=True)
    966        jbytes = j.encode("utf-8")
    967        start = 8 + len(jbytes)
    968        with open(path, mode="wb") as f:  # make empty file, fill with enough space
    969            f.write(b"\x00" * (start + offset))
    970        with open(path, mode="r+b") as f:
    971            with mmap.mmap(f.fileno(), length=0, access=mmap.ACCESS_WRITE) as m:
    972                m[0:8] = np.int64(len(j)).tobytes()
    973                m[8:start] = jbytes
    974                for t, tm in zip(tensors.values(), metadata.values()):
    975                    data_start, data_end = tm["data_offsets"]
    976                    m[start + data_start : start + data_end] = t.numpy().tobytes()
Backend( operator_set: OperatorSet, procedure_set: ProcedureSet) View Source

    797    def __init__(
    798        self,
    799        operator_set: OperatorSet,
    800        procedure_set: ProcedureSet,
    801    ):
    802        self.operator_set = operator_set
    803        self.procedure_set = procedure_set
    804        self.node_types = dict()
    805        self.impls = dict()
    806        self.register_node(tuple, lambda t: (None, t), lambda _, xs: tuple(xs), "tuple")
    807        self.register_node(list, lambda l: (None, l), lambda _, xs: list(xs), "list")
    808        self.register_node(
    809            dict,
    810            lambda d: list_map(tuple, unzip2(sorted(d.items()))),
    811            lambda keys, vals: dict(list_zip(keys, vals)),
    812            "dict",
    813        )
    814        self.register_node(
    815            UndefinedPrimal,
    816            lambda u: (u.symval, ()),
    817            lambda symval, _: UndefinedPrimal(symval),
    818            "UndefinedPrimal",
    819        )
LOG_LRU = 0
LOG_JIT = 0
LOG_TREE = 0
LOG_BACKEND = 0
LOG_PROGRAM = 0
LOG_INIT = 1
device_var = 'cpu:0'
DEFAULT_DEVICE = <Device: 'cpu:0'>
DEFAULT_DTYPE = <DType: float32>
dtype_for_indices: DType = None
operator_set
procedure_set
node_types
impls
def set_impl(self, op: function): View Source

    821    def set_impl(self, op: Union[types.LambdaType, types.FunctionType]):
    822        def set_impl_(fn):
    823            self.impls[op] = types.MethodType(fn, self)
    824
    825        return set_impl_
def register_node( self, ty: Type, to_iter: Callable, from_iter: Callable, name=None) -> None: View Source

    827    def register_node(self, ty: Type, to_iter: Callable, from_iter: Callable, name=None) -> None:
    828        if name is None:
    829            name = str(ty)
    830        self.node_types[ty] = NodeType(name, to_iter, from_iter)
def tensor( self, val: Union[list, tuple, numpy.ndarray, TensorBuffer] = None, dtype: Optional[Any] = None, device=None): View Source

    855    def tensor(
    856        self,
    857        val: Union[list, tuple, np.ndarray, "TensorBuffer"] = None,
    858        dtype: Optional[Any] = None,
    859        device=None,
    860    ):
    861        if isinstance(val, TensorBuffer):
    862            return Tensor(val)
    863        elif isinstance(val, Tensor):
    864            return val
    865        if type(val) is bytes:
    866            val = np.frombuffer(val, dtype=dtype)
    867        return self.from_numpy(val, dtype, device)
def symbolic_tensor( self, shape: Union[list, tuple, numpy.ndarray, TensorBuffer] = None, dtype: Optional[Any] = None, device=None): View Source

    869    def symbolic_tensor(
    870        self,
    871        shape: Union[list, tuple, np.ndarray, "TensorBuffer"] = None,
    872        dtype: Optional[Any] = None,
    873        device=None,
    874    ):
    875        dtype = dtype or self.DEFAULT_DTYPE
    876        device = device or self.DEFAULT_DEVICE
    877        return SymbolicTensor(shape, dtype, device)
def seed(self, seed): View Source

    879    def seed(self, seed):
    880        raise NotImplementedError
default_dtype_value View Source

    882    @property
    883    def default_dtype_value(self):
    884        return self.dtype_map[backend.DEFAULT_DTYPE]
def set_method(self, method): View Source

    886    def set_method(self, method):
    887        setattr(self, method.__name__, types.MethodType(method, self))
def from_numpy(self, val, device): View Source

    889    def from_numpy(self, val, device):
    890        raise NotImplementedError
def numpy_of(self, tensor): View Source

    892    def numpy_of(self, tensor):
    893        raise NotImplementedError
def device_of(self, tensor): View Source

    895    def device_of(self, tensor):
    896        raise NotImplementedError
def shape_of(self, tensor): View Source

    898    def shape_of(self, tensor):
    899        raise NotImplementedError
def dtype_of(self, tensor): View Source

    901    def dtype_of(self, tensor):
    902        raise NotImplementedError
def jit_program(*args, **kwargs) -> Any: View Source

    172        def decorated_function(*args, **kwargs) -> Any:
    173            result = wrapper(*args, **kwargs)
    174            cache_info = wrapper.cache_info()
    175
    176            dblog(
    177                f"{fn.__name__}.{cache_info} {args.__hash__()}",
    178                enable=backend.LOG_LRU,
    179            )
    180            tb = "".join(traceback.format_list(traceback.extract_stack())[tb_start:tb_end]).replace("\n    ", ":\t") + "-" * 20 + "\n"
    181            dblog(f"{tb}", enable=backend.LOG_LRU)
    182
    183            return result
def codegen( self, program: Program, args: Tuple, in_symvals: Tuple, name: str): View Source

    919    def codegen(self, program: "Program", args: Tuple, in_symvals: Tuple, name: str):
    920        "Returns compiler IR from the Program"
    921        raise NotImplementedError

Returns compiler IR from the Program

def compile( self, program: Program, args: Tuple, in_symvals: Tuple, name: str): View Source

    923    def compile(self, program: "Program", args: Tuple, in_symvals: Tuple, name: str):
    924        "Compiles compiler IR to a Python callable function"
    925        raise NotImplementedError

Compiles compiler IR to a Python callable function

def export(self, jit_output, *args, **params): View Source

    927    def export(self, jit_output, *args, **params):
    928        raise NotImplementedError
def load(self, path, single_key='_tensor'): View Source

    930    def load(self, path, single_key="_tensor"):
    931        with open(path, mode="rb") as f:
    932            with mmap.mmap(f.fileno(), length=0, access=mmap.ACCESS_READ) as m:
    933                json_len = np.int64(m[0])
    934                start = 8 + json_len
    935                metadata = json.loads(m[8:start])
    936                ret = {}
    937                for k, v in metadata.items():
    938                    if k != "__metadata__":
    939                        dtype = Tensor.mlir_dtype_map[(v["dtype"])]
    940                        data_start = start + v["data_offsets"][0]
    941                        data_end = start + v["data_offsets"][1]
    942                        t_np = np.frombuffer(m[data_start:data_end], dtype=dtype.numpy())
    943                        t = backend.tensor(t_np, dtype=dtype)
    944                        t = t.reshape(tuple(v["shape"]))
    945                        ret[k] = t
    946                if len(ret) == 1 and single_key in ret.keys():
    947                    return ret[single_key]
    948                return ret
def save( self, tensors: Dict[str, Tensor], path: str, single_key='_tensor'): View Source

    950    def save(self, tensors: Dict[str, Tensor], path: str, single_key="_tensor"):
    951        if isinstance(tensors, Tensor):
    952            tensors = {single_key: tensors}
    953        else:
    954            assert all((isinstance(k, str) and isinstance(v, Tensor)) for k, v in tensors.items())
    955
    956        metadata, offset = {}, 0
    957        for k, v in tensors.items():
    958            metadata[k] = {
    959                "dtype": v.dtype.mlir,
    960                "shape": list(v.shape),
    961                "data_offsets": [offset, offset + v.nbytes()],
    962            }
    963            offset += v.nbytes()
    964        j = json.dumps(metadata, separators=(",", ":"))
    965        Path(path).unlink(missing_ok=True)
    966        jbytes = j.encode("utf-8")
    967        start = 8 + len(jbytes)
    968        with open(path, mode="wb") as f:  # make empty file, fill with enough space
    969            f.write(b"\x00" * (start + offset))
    970        with open(path, mode="r+b") as f:
    971            with mmap.mmap(f.fileno(), length=0, access=mmap.ACCESS_WRITE) as m:
    972                m[0:8] = np.int64(len(j)).tobytes()
    973                m[8:start] = jbytes
    974                for t, tm in zip(tensors.values(), metadata.values()):
    975                    data_start, data_end = tm["data_offsets"]
    976                    m[start + data_start : start + data_end] = t.numpy().tobytes()
class Var: View Source

    984class Var:
    985    def __init__(self, symval):
    986        self.symval = symval
    987        self.val = None
Var(symval) View Source

    985    def __init__(self, symval):
    986        self.symval = symval
    987        self.val = None
symval
val
class Lit: View Source

    990class Lit:
    991    def __init__(self, val):
    992        self.symval = SymbolicTensor.like(get_symval(val))
    993        self.val = val
Lit(val) View Source

    991    def __init__(self, val):
    992        self.symval = SymbolicTensor.like(get_symval(val))
    993        self.val = val
symval
val
Atom = typing.Union[Var, Lit]
class Instruction(typing.NamedTuple): View Source

     999class Instruction(NamedTuple):
    1000    op: Operator
    1001    inputs: List[Atom]
    1002    params: Dict[str, Any]
    1003    out_binders: List[Atom]

Instruction(op, inputs, params, out_binders)

Instruction( op: Operator, inputs: List[Union[Var, Lit]], params: Dict[str, Any], out_binders: List[Union[Var, Lit]])

Create new instance of Instruction(op, inputs, params, out_binders)

op: Operator

Alias for field number 0

inputs: List[Union[Var, Lit]]

Alias for field number 1

params: Dict[str, Any]

Alias for field number 2

out_binders: List[Union[Var, Lit]]

Alias for field number 3

##### Inherited Members

builtins.tuple
  index
  count

class ProgramEnvVar(typing.NamedTuple): View Source

    1006class ProgramEnvVar(NamedTuple):
    1007    name: str
    1008    symval: SymbolicTensor
    1009    is_const: bool = False
    1010
    1011    @property
    1012    def shape(self):
    1013        return self.symval.shape
    1014
    1015    @property
    1016    def dtype(self):
    1017        return self.symval.dtype
    1018
    1019    @property
    1020    def device(self):
    1021        return self.symval.device
    1022
    1023    @property
    1024    def ndim(self):
    1025        return self.symval.ndim
    1026
    1027    def numpy(self):
    1028        return self.symval.numpy()
    1029
    1030    def __repr__(self):
    1031        return f"<ProgramEnvVar: name={self.name}, symval={self.symval}>"
    1032
    1033    str_short = __repr__

ProgramEnvVar(name, symval, is_const)

ProgramEnvVar(name: str, symval: SymbolicTensor, is_const: bool = False)

Create new instance of ProgramEnvVar(name, symval, is_const)

name: str

Alias for field number 0

symval: SymbolicTensor

Alias for field number 1

is_const: bool

Alias for field number 2

shape View Source

    1011    @property
    1012    def shape(self):
    1013        return self.symval.shape
dtype View Source

    1015    @property
    1016    def dtype(self):
    1017        return self.symval.dtype
device View Source

    1019    @property
    1020    def device(self):
    1021        return self.symval.device
ndim View Source

    1023    @property
    1024    def ndim(self):
    1025        return self.symval.ndim
def numpy(self): View Source

    1027    def numpy(self):
    1028        return self.symval.numpy()
def str_short(self): View Source

    1030    def __repr__(self):
    1031        return f"<ProgramEnvVar: name={self.name}, symval={self.symval}>"

Return repr(self).

##### Inherited Members

builtins.tuple
  index
  count

class Program: View Source

    1036class Program:
    1037    def __init__(
    1038        self,
    1039        in_binders: Any,
    1040        instructions: Tuple[Instruction],
    1041        outs: Any,
    1042        num_consts: int = 0,
    1043        static_args: Any = (),
    1044        name: str = "my_program",
    1045        indent_amount=4,
    1046    ):
    1047        self.in_binders: Any = in_binders
    1048        self.outs: Any = outs
    1049        self.instructions = self.prune_instructions(instructions, outs)
    1050        # self.instructions = instructions
    1051        self.num_consts: int = num_consts
    1052        self.static_args = static_args
    1053        self.name: str = name
    1054        self.indent_amount: int = indent_amount
    1055
    1056        self.env: Dict[ProgramEnvVar, Any] = dict()
    1057        for inb in self.in_binders:
    1058            prefix = "x" if type(inb.symval) is SymbolicTensor else "c"
    1059            idx = sum([1 if v.name[0] == prefix else 0 for v in self.env.values()])
    1060            self.env[inb] = ProgramEnvVar(f"{prefix}{idx}", inb.symval, True if prefix == "c" else False)
    1061        for instruction in self.instructions:
    1062            if len(instruction.out_binders) == 0:
    1063                continue
    1064            for outb in instruction.out_binders:
    1065                prefix = "y" if outb in self.outs else "z"
    1066                idx = sum([1 if v.name[0] == prefix else 0 for v in self.env.values()])
    1067                self.env[outb] = ProgramEnvVar(f"{prefix}{idx}", outb.symval)
    1068        self.curr_repr = repr(self)
    1069
    1070    def pprint_shape(self, symval, scalar_as_empty_array=False):
    1071        xdtype = symval.dtype.mlir
    1072        if len(symval.shape) > 0:
    1073            xshape = f"{', '.join((repr(i) for i in symval.shape))}"
    1074            return f"[{xshape}, {xdtype}]"
    1075        else:
    1076            return f"[{xdtype}]"
    1077
    1078    def pprint_sig(self, in_symvals, out_symvals, unpack_unary_output=False):
    1079        in_code = ", ".join(self.pprint_shape(t) for t in in_symvals)
    1080        in_code = f"({in_code})" if len(in_symvals) > 1 else in_code
    1081        out_code = ", ".join(self.pprint_shape(t) for t in out_symvals)
    1082        out_code = f"({out_code})" if len(out_symvals) > 1 or unpack_unary_output else out_code
    1083        typing_code = f"{in_code} -> {out_code}"
    1084        return typing_code
    1085
    1086    def __repr__(self):
    1087        fn_defs = self.instructions_as_code(self, dict())
    1088        return "\n".join(line for code_lines in fn_defs.values() for line in code_lines)
    1089
    1090    def save(self, *args, dir_path="/tmp/slope_program", dry_run=False):
    1091        os.makedirs(dir_path, exist_ok=True)
    1092        head_code_lines = [f"import slope # backend={backend.__class__.__name__}"]
    1093        fn_defs = self.instructions_as_code(self, dict())
    1094        in_binders_vars = [self.env[i] for i in self.in_binders]
    1095        for i in range(len(self.in_binders)):
    1096            ibv = in_binders_vars[i]
    1097            if ibv.is_const:
    1098                const_filename = f"{ibv.name}.safetensors"
    1099                const_path = os.path.join(dir_path, f"{const_filename}")
    1100                if not dry_run:
    1101                    backend.save(args[i], const_path)
    1102                dblog(
    1103                    f"Saved {ibv.name} at {const_path}",
    1104                    enable=backend.LOG_BACKEND,
    1105                )
    1106                head_code_lines += [f"""{ibv.name} = slope.load("./{const_filename}")"""]
    1107        head_code_lines += [""]
    1108        code = "\n".join(head_code_lines + [line for code_lines in fn_defs.values() for line in code_lines])
    1109        dblog(
    1110            f"Contents of {self.name}:\n```\n{code}\n```",
    1111            enable=backend.LOG_BACKEND,
    1112        )
    1113        program_path = os.path.join(dir_path, "main.py")
    1114        if not dry_run:
    1115            with open(program_path, "w") as f:
    1116                f.write(code)
    1117        dblog(
    1118            f"Saved program {self.name} at {program_path}",
    1119            enable=backend.LOG_BACKEND,
    1120        )
    1121        ls_contents = "\n\t".join(os.listdir(dir_path))
    1122        dblog(
    1123            f"Contents of {dir_path}:\n\t{ls_contents}",
    1124            enable=backend.LOG_BACKEND,
    1125        )
    1126
    1127    def __hash__(self):
    1128        return hash(self.curr_repr)
    1129
    1130    def __eq__(self, other):
    1131        return self is other
    1132
    1133    @classmethod
    1134    def instructions_as_code(cls, program, fn_defs):
    1135        def indent(code, indent_amount):
    1136            spaces = " " * (len(code) - len(code.lstrip()))
    1137            spaces += " " * indent_amount
    1138            return "\n".join([spaces + line for line in code.strip().split("\n")])
    1139
    1140        in_binders_vars = [program.env[i] for i in program.in_binders]
    1141        body_code_lines = []
    1142        for instruction in program.instructions:
    1143            if len(instruction.out_binders) == 0:
    1144                continue
    1145            params = instruction.params.copy()
    1146            for param_name, param in params.items():
    1147                if isinstance(param, Program):
    1148                    sub_program = param
    1149                    fn_defs = cls.instructions_as_code(sub_program, fn_defs)
    1150                    program_in_vals = ", ".join(f"{program.env[x].name}" for x in instruction.inputs)
    1151                    params[param_name] = f"slope.make_program({sub_program.name}, {program_in_vals})[0]"
    1152                if isinstance(param, DType):
    1153                    params[param_name] = f"slope.{param.name}"
    1154            param_vals = ", ".join(f"{param_name}={param}" for param_name, param in params.items())
    1155            in_vals = ", ".join(f"{program.env[x].name}" for x in instruction.inputs)
    1156            out_vals = ", ".join(f"{program.env[z].name}" for z in instruction.out_binders)
    1157            sig = program.pprint_sig(
    1158                [program.env[x].symval for x in instruction.inputs],
    1159                [program.env[y].symval for y in instruction.out_binders],
    1160            )
    1161            line = f"""{out_vals} = slope.{instruction.op.name}({in_vals}{", " if (param_vals and in_vals) else ""}{param_vals}) # {sig}"""
    1162            body_code_lines += [indent(line, program.indent_amount)]
    1163
    1164        fn_args_str = ", ".join([f"{i.name}" for i in in_binders_vars])
    1165        # fn_static_args_str = ", ".join([f"{a}={a_val}" for a, a_val in program.static_args])
    1166        out_vars = [program.env[o] for o in program.outs]
    1167        fn_sig = program.pprint_sig(
    1168            [i.symval for i in in_binders_vars],
    1169            [o.symval for o in out_vars],
    1170        )
    1171        head_code_line = [f"def {program.name}({fn_args_str}): # {fn_sig}"]
    1172        out_str = ", ".join([f"{o.name}" for o in out_vars])
    1173        tail_code_line = [indent(f"return {out_str}", program.indent_amount)]
    1174        code_lines = head_code_line + body_code_lines + tail_code_line + ["\n"]
    1175
    1176        fn_defs[program.name] = code_lines
    1177        return fn_defs
    1178
    1179    @staticmethod
    1180    def prune_instructions(instructions, outs):
    1181        graph = dict()
    1182        for instruction in instructions:
    1183            parent_nodes, child_nodes = instruction.out_binders, instruction.inputs
    1184            for parent in parent_nodes:
    1185                if parent not in graph:
    1186                    graph[parent] = set()
    1187                for child in child_nodes:
    1188                    graph[parent].add(child)
    1189        visited_from_terminal = set()
    1190
    1191        def dfs(node, visited):
    1192            visited.add(node)
    1193            if node in graph:
    1194                for neighbor in graph[node]:
    1195                    if neighbor not in visited:
    1196                        dfs(neighbor, visited)
    1197
    1198        for terminal_node in outs:
    1199            dfs(terminal_node, visited_from_terminal)
    1200        unreachable_nodes = set(graph.keys()) - visited_from_terminal
    1201
    1202        instructions_to_prune = []
    1203        for instruction in instructions:
    1204            parent_nodes, child_nodes = instruction.out_binders, instruction.inputs
    1205            if any(node in unreachable_nodes for node in parent_nodes) or any(node in unreachable_nodes for node in child_nodes):
    1206                instructions_to_prune += [instruction]
    1207        new_instructions = [inst for inst in instructions if inst not in instructions_to_prune]
    1208        if backend.LOG_PROGRAM:
    1209            LI = len(instructions)
    1210            LNI = len(new_instructions)
    1211            DIFF = LI - LNI
    1212            UN = len(unreachable_nodes)
    1213            dblog(f"Before: {LI}\tAfter: {LNI}\tDiff vs Unreachables: {DIFF} == {UN} = {DIFF==UN}")
    1214        return new_instructions
Program( in_binders: Any, instructions: Tuple[Instruction], outs: Any, num_consts: int = 0, static_args: Any = (), name: str = 'my_program', indent_amount=4) View Source

    1037    def __init__(
    1038        self,
    1039        in_binders: Any,
    1040        instructions: Tuple[Instruction],
    1041        outs: Any,
    1042        num_consts: int = 0,
    1043        static_args: Any = (),
    1044        name: str = "my_program",
    1045        indent_amount=4,
    1046    ):
    1047        self.in_binders: Any = in_binders
    1048        self.outs: Any = outs
    1049        self.instructions = self.prune_instructions(instructions, outs)
    1050        # self.instructions = instructions
    1051        self.num_consts: int = num_consts
    1052        self.static_args = static_args
    1053        self.name: str = name
    1054        self.indent_amount: int = indent_amount
    1055
    1056        self.env: Dict[ProgramEnvVar, Any] = dict()
    1057        for inb in self.in_binders:
    1058            prefix = "x" if type(inb.symval) is SymbolicTensor else "c"
    1059            idx = sum([1 if v.name[0] == prefix else 0 for v in self.env.values()])
    1060            self.env[inb] = ProgramEnvVar(f"{prefix}{idx}", inb.symval, True if prefix == "c" else False)
    1061        for instruction in self.instructions:
    1062            if len(instruction.out_binders) == 0:
    1063                continue
    1064            for outb in instruction.out_binders:
    1065                prefix = "y" if outb in self.outs else "z"
    1066                idx = sum([1 if v.name[0] == prefix else 0 for v in self.env.values()])
    1067                self.env[outb] = ProgramEnvVar(f"{prefix}{idx}", outb.symval)
    1068        self.curr_repr = repr(self)
in_binders: Any
outs: Any
instructions
num_consts: int
static_args
name: str
indent_amount: int
env: Dict[ProgramEnvVar, Any]
curr_repr
def pprint_shape(self, symval, scalar_as_empty_array=False): View Source

    1070    def pprint_shape(self, symval, scalar_as_empty_array=False):
    1071        xdtype = symval.dtype.mlir
    1072        if len(symval.shape) > 0:
    1073            xshape = f"{', '.join((repr(i) for i in symval.shape))}"
    1074            return f"[{xshape}, {xdtype}]"
    1075        else:
    1076            return f"[{xdtype}]"
def pprint_sig(self, in_symvals, out_symvals, unpack_unary_output=False): View Source

    1078    def pprint_sig(self, in_symvals, out_symvals, unpack_unary_output=False):
    1079        in_code = ", ".join(self.pprint_shape(t) for t in in_symvals)
    1080        in_code = f"({in_code})" if len(in_symvals) > 1 else in_code
    1081        out_code = ", ".join(self.pprint_shape(t) for t in out_symvals)
    1082        out_code = f"({out_code})" if len(out_symvals) > 1 or unpack_unary_output else out_code
    1083        typing_code = f"{in_code} -> {out_code}"
    1084        return typing_code
def save(self, *args, dir_path='/tmp/slope_program', dry_run=False): View Source

    1090    def save(self, *args, dir_path="/tmp/slope_program", dry_run=False):
    1091        os.makedirs(dir_path, exist_ok=True)
    1092        head_code_lines = [f"import slope # backend={backend.__class__.__name__}"]
    1093        fn_defs = self.instructions_as_code(self, dict())
    1094        in_binders_vars = [self.env[i] for i in self.in_binders]
    1095        for i in range(len(self.in_binders)):
    1096            ibv = in_binders_vars[i]
    1097            if ibv.is_const:
    1098                const_filename = f"{ibv.name}.safetensors"
    1099                const_path = os.path.join(dir_path, f"{const_filename}")
    1100                if not dry_run:
    1101                    backend.save(args[i], const_path)
    1102                dblog(
    1103                    f"Saved {ibv.name} at {const_path}",
    1104                    enable=backend.LOG_BACKEND,
    1105                )
    1106                head_code_lines += [f"""{ibv.name} = slope.load("./{const_filename}")"""]
    1107        head_code_lines += [""]
    1108        code = "\n".join(head_code_lines + [line for code_lines in fn_defs.values() for line in code_lines])
    1109        dblog(
    1110            f"Contents of {self.name}:\n```\n{code}\n```",
    1111            enable=backend.LOG_BACKEND,
    1112        )
    1113        program_path = os.path.join(dir_path, "main.py")
    1114        if not dry_run:
    1115            with open(program_path, "w") as f:
    1116                f.write(code)
    1117        dblog(
    1118            f"Saved program {self.name} at {program_path}",
    1119            enable=backend.LOG_BACKEND,
    1120        )
    1121        ls_contents = "\n\t".join(os.listdir(dir_path))
    1122        dblog(
    1123            f"Contents of {dir_path}:\n\t{ls_contents}",
    1124            enable=backend.LOG_BACKEND,
    1125        )
@classmethod
def instructions_as_code(cls, program, fn_defs): View Source

    1133    @classmethod
    1134    def instructions_as_code(cls, program, fn_defs):
    1135        def indent(code, indent_amount):
    1136            spaces = " " * (len(code) - len(code.lstrip()))
    1137            spaces += " " * indent_amount
    1138            return "\n".join([spaces + line for line in code.strip().split("\n")])
    1139
    1140        in_binders_vars = [program.env[i] for i in program.in_binders]
    1141        body_code_lines = []
    1142        for instruction in program.instructions:
    1143            if len(instruction.out_binders) == 0:
    1144                continue
    1145            params = instruction.params.copy()
    1146            for param_name, param in params.items():
    1147                if isinstance(param, Program):
    1148                    sub_program = param
    1149                    fn_defs = cls.instructions_as_code(sub_program, fn_defs)
    1150                    program_in_vals = ", ".join(f"{program.env[x].name}" for x in instruction.inputs)
    1151                    params[param_name] = f"slope.make_program({sub_program.name}, {program_in_vals})[0]"
    1152                if isinstance(param, DType):
    1153                    params[param_name] = f"slope.{param.name}"
    1154            param_vals = ", ".join(f"{param_name}={param}" for param_name, param in params.items())
    1155            in_vals = ", ".join(f"{program.env[x].name}" for x in instruction.inputs)
    1156            out_vals = ", ".join(f"{program.env[z].name}" for z in instruction.out_binders)
    1157            sig = program.pprint_sig(
    1158                [program.env[x].symval for x in instruction.inputs],
    1159                [program.env[y].symval for y in instruction.out_binders],
    1160            )
    1161            line = f"""{out_vals} = slope.{instruction.op.name}({in_vals}{", " if (param_vals and in_vals) else ""}{param_vals}) # {sig}"""
    1162            body_code_lines += [indent(line, program.indent_amount)]
    1163
    1164        fn_args_str = ", ".join([f"{i.name}" for i in in_binders_vars])
    1165        # fn_static_args_str = ", ".join([f"{a}={a_val}" for a, a_val in program.static_args])
    1166        out_vars = [program.env[o] for o in program.outs]
    1167        fn_sig = program.pprint_sig(
    1168            [i.symval for i in in_binders_vars],
    1169            [o.symval for o in out_vars],
    1170        )
    1171        head_code_line = [f"def {program.name}({fn_args_str}): # {fn_sig}"]
    1172        out_str = ", ".join([f"{o.name}" for o in out_vars])
    1173        tail_code_line = [indent(f"return {out_str}", program.indent_amount)]
    1174        code_lines = head_code_line + body_code_lines + tail_code_line + ["\n"]
    1175
    1176        fn_defs[program.name] = code_lines
    1177        return fn_defs
@staticmethod
def prune_instructions(instructions, outs): View Source

    1179    @staticmethod
    1180    def prune_instructions(instructions, outs):
    1181        graph = dict()
    1182        for instruction in instructions:
    1183            parent_nodes, child_nodes = instruction.out_binders, instruction.inputs
    1184            for parent in parent_nodes:
    1185                if parent not in graph:
    1186                    graph[parent] = set()
    1187                for child in child_nodes:
    1188                    graph[parent].add(child)
    1189        visited_from_terminal = set()
    1190
    1191        def dfs(node, visited):
    1192            visited.add(node)
    1193            if node in graph:
    1194                for neighbor in graph[node]:
    1195                    if neighbor not in visited:
    1196                        dfs(neighbor, visited)
    1197
    1198        for terminal_node in outs:
    1199            dfs(terminal_node, visited_from_terminal)
    1200        unreachable_nodes = set(graph.keys()) - visited_from_terminal
    1201
    1202        instructions_to_prune = []
    1203        for instruction in instructions:
    1204            parent_nodes, child_nodes = instruction.out_binders, instruction.inputs
    1205            if any(node in unreachable_nodes for node in parent_nodes) or any(node in unreachable_nodes for node in child_nodes):
    1206                instructions_to_prune += [instruction]
    1207        new_instructions = [inst for inst in instructions if inst not in instructions_to_prune]
    1208        if backend.LOG_PROGRAM:
    1209            LI = len(instructions)
    1210            LNI = len(new_instructions)
    1211            DIFF = LI - LNI
    1212            UN = len(unreachable_nodes)
    1213            dblog(f"Before: {LI}\tAfter: {LNI}\tDiff vs Unreachables: {DIFF} == {UN} = {DIFF==UN}")
    1214        return new_instructions
class ProgramType(typing.NamedTuple): View Source

    1217class ProgramType(NamedTuple):
    1218    in_types: Tuple[SymbolicTensor]
    1219    out_types: Tuple[SymbolicTensor]
    1220
    1221    def __repr__(self):
    1222        in_types = ", ".join(symval.str_short() for symval in self.in_types)
    1223        out_types = ", ".join(symval.str_short() for symval in self.out_types)
    1224        return f"({in_types}) -> ({out_types})"

ProgramType(in_types, out_types)

ProgramType( in_types: Tuple[SymbolicTensor], out_types: Tuple[SymbolicTensor])

Create new instance of ProgramType(in_types, out_types)

in_types: Tuple[SymbolicTensor]

Alias for field number 0

out_types: Tuple[SymbolicTensor]

Alias for field number 1

##### Inherited Members

builtins.tuple
  index
  count

class Empty: View Source

    1232class Empty:
    1233    pass
empty = <Empty object>
class Store: View Source

    1239class Store:
    1240    val = empty
    1241
    1242    def set_value(self, val):
    1243        assert self.val is empty
    1244        self.val = val
    1245
    1246    def __call__(self):
    1247        return self.val
val = <Empty object>
def set_value(self, val): View Source

    1242    def set_value(self, val):
    1243        assert self.val is empty
    1244        self.val = val
class NodeType(typing.NamedTuple): View Source

    1250class NodeType(NamedTuple):
    1251    name: str
    1252    flatten: Callable
    1253    unflatten: Callable

NodeType(name, flatten, unflatten)

NodeType(name: str, flatten: Callable, unflatten: Callable)

Create new instance of NodeType(name, flatten, unflatten)

name: str

Alias for field number 0

flatten: Callable

Alias for field number 1

unflatten: Callable

Alias for field number 2

##### Inherited Members

builtins.tuple
  index
  count

class TreeDef(typing.NamedTuple): View Source

    1256class TreeDef(NamedTuple):
    1257    node_type: NodeType
    1258    node_metadata: Hashable
    1259    child_treedefs: Tuple["TreeDef", ...]
    1260
    1261    def __repr__(self):
    1262        ret = self.tree_repr(self)
    1263        return ret
    1264
    1265    def tree_repr(self, tree, indent="  ", prefix="", last=True):
    1266        ret = ""
    1267
    1268        def _tree_repr(tree, indent, prefix, last):
    1269            nonlocal ret
    1270            if isinstance(tree, TreeDef):
    1271                ret += f'{prefix} {("└─" if last else "├─")} {tree.node_type.name}\n'
    1272                for i, item in enumerate(tree.child_treedefs):
    1273                    new_prefix = prefix + (indent if not last else "   ")
    1274                    new_last = i == len(tree.child_treedefs) - 1
    1275                    _tree_repr(item, indent, new_prefix, new_last)
    1276            else:
    1277                ret += f'{prefix} {("└─" if last else "├─")} {tree}\n'
    1278
    1279        _tree_repr(tree, indent="  ", prefix="", last=True)
    1280        return ret
    1281
    1282    @property
    1283    def num_leaves(self):
    1284        def get_num_leaves(x):
    1285            if isinstance(x, Leaf):
    1286                return 1
    1287            else:
    1288                return sum(get_num_leaves(sub_x) for sub_x in x.child_treedefs)
    1289
    1290        return sum(get_num_leaves(x) for x in self.child_treedefs)

TreeDef(node_type, node_metadata, child_treedefs)

TreeDef( node_type: NodeType, node_metadata: Hashable, child_treedefs: Tuple[ForwardRef('TreeDef'), ...])

Create new instance of TreeDef(node_type, node_metadata, child_treedefs)

node_type: NodeType

Alias for field number 0

node_metadata: Hashable

Alias for field number 1

child_treedefs: Tuple[TreeDef, ...]

Alias for field number 2

def tree_repr(self, tree, indent=' ', prefix='', last=True): View Source

    1265    def tree_repr(self, tree, indent="  ", prefix="", last=True):
    1266        ret = ""
    1267
    1268        def _tree_repr(tree, indent, prefix, last):
    1269            nonlocal ret
    1270            if isinstance(tree, TreeDef):
    1271                ret += f'{prefix} {("└─" if last else "├─")} {tree.node_type.name}\n'
    1272                for i, item in enumerate(tree.child_treedefs):
    1273                    new_prefix = prefix + (indent if not last else "   ")
    1274                    new_last = i == len(tree.child_treedefs) - 1
    1275                    _tree_repr(item, indent, new_prefix, new_last)
    1276            else:
    1277                ret += f'{prefix} {("└─" if last else "├─")} {tree}\n'
    1278
    1279        _tree_repr(tree, indent="  ", prefix="", last=True)
    1280        return ret
num_leaves View Source

    1282    @property
    1283    def num_leaves(self):
    1284        def get_num_leaves(x):
    1285            if isinstance(x, Leaf):
    1286                return 1
    1287            else:
    1288                return sum(get_num_leaves(sub_x) for sub_x in x.child_treedefs)
    1289
    1290        return sum(get_num_leaves(x) for x in self.child_treedefs)

##### Inherited Members

builtins.tuple
  index
  count

class Leaf: View Source

    1293class Leaf:
    1294    def __init__(self, val):
    1295        if hasattr(val, "shape"):
    1296            val = SymbolicTensor.like(val)
    1297        self.val = val
    1298
    1299    def __repr__(self):
    1300        ret = self.val.str_short() if isinstance(self.val, SymbolicTensor) else repr(self.val)
    1301        return f"<Leaf: {ret}>"
    1302
    1303    def __hash__(self):
    1304        return hash(self.val)
    1305
    1306    def __eq__(self, other):
    1307        return True  # make TreeDef __eq__ don't care Leaf
    1308        # if isinstance(other, Leaf): # TODO: test above assumption
    1309        #     return self.val == other.val
Leaf(val) View Source

    1294    def __init__(self, val):
    1295        if hasattr(val, "shape"):
    1296            val = SymbolicTensor.like(val)
    1297        self.val = val
val
class JitOutput: View Source

    1317class JitOutput:
    1318    def __init__(self, program: Program, codegen_output: CodegenOutput, fn, code: str, consts: List[Any]):
    1319        super().__init__()
    1320        self.program = program
    1321        self.code = code
    1322        self.codegen_output = codegen_output
    1323        self.fn: Callable = fn
    1324        self.consts = consts
    1325
    1326    def __call__(self, *args, **params):
    1327        args, in_tree = tree_flatten(args)
    1328        args = tree_map(lambda a: a.val if isinstance(a, Tensor) else a, args)
    1329        try:
    1330            outs = self.fn(*args, **params)
    1331            if not isinstance(outs, tuple):  # TODO: IREE FunctionInvoker destructure 1-tuple, need to undo
    1332                outs = (outs,)
    1333        except Exception as e:
    1334            dblog(self.code, enable=backend.LOG_JIT)
    1335            raise
    1336        return [backend.tensor(TensorBuffer(o)) for o in outs]
JitOutput( program: Program, codegen_output: CodegenOutput, fn, code: str, consts: List[Any]) View Source

    1318    def __init__(self, program: Program, codegen_output: CodegenOutput, fn, code: str, consts: List[Any]):
    1319        super().__init__()
    1320        self.program = program
    1321        self.code = code
    1322        self.codegen_output = codegen_output
    1323        self.fn: Callable = fn
    1324        self.consts = consts
program
code
codegen_output
fn: Callable
consts
class JitOp(MetaOperator): View Source

    1339class JitOp(MetaOperator):
    1340    def meta_impl(self, *args, program: Program, **_):
    1341        hashed_program = Hashed(program)
    1342        num_consts = program.num_consts
    1343        consts, args = args[:num_consts], args[num_consts:]
    1344        hashed_consts = tuple(map(Hashed, consts))
    1345        jit_output = backend.jit_program(hashed_program, hashed_consts)
    1346        ret = jit_output(*consts, *args)
    1347        return ret
    1348
    1349    def reorg_args(self, args, params):
    1350        return args, params
    1351
    1352    def typecheck(self, *in_types, program: Program):
    1353        program_type = typecheck_program(program)
    1354        if not all(t1 == t2 for t1, t2 in zip(program_type.in_types, in_types)):
    1355            ret = "Type mismatch program.in_types vs in_types:\n"
    1356            for i, j in zip(program_type.in_types, in_types):
    1357                ret += f"{i}, {j}, {i == j}"
    1358            raise TypeError(ret)
    1359        return program_type.out_types
    1360
    1361    def vmap(self, dim_size, vals_in, dims_in, program: Program):
    1362        program, consts = vmap_program(program, dim_size, tuple(dims_in))
    1363        outs = self(*consts, *vals_in, program=program)
    1364        if not isinstance(outs, tuple):
    1365            outs = (outs,)
    1366        return outs, [0] * len(outs)
    1367
    1368    def jvp(self, primals, tangents, *, program):
    1369        new_program, new_consts = jvp_program(program)
    1370        outs = bind(
    1371            self,
    1372            *new_consts,
    1373            *primals,
    1374            *tangents,
    1375            program=new_program,
    1376        )
    1377        n = len(outs) // 2
    1378        primals_out, tangents_out = outs[:n], outs[n:]
    1379        return primals_out, tangents_out
    1380
    1381    def T(self, cotangents, *invals, program):
    1382        undef_primals = [isinstance(x, UndefinedPrimal) for x in invals]
    1383        transposed_program, new_consts = transpose_program(program, tuple(undef_primals))
    1384
    1385        residuals, _ = partition_list(undef_primals, invals)
    1386        outs = bind(
    1387            self,
    1388            *new_consts,
    1389            *residuals,
    1390            *cotangents,
    1391            program=transposed_program,
    1392        )
    1393        outs = iter(outs)
    1394
    1395        return [next(outs) if undef else None for undef in undef_primals]
    1396
    1397    def partial_run(self, trace, tracers, *, program):
    1398        in_unknowns = [not t.pval.is_known for t in tracers]
    1399        program1, program2, out_unknowns, num_res = partial_run_program(program, in_unknowns)
    1400        known_tracers, unknown_tracers = partition_list(in_unknowns, tracers)
    1401        known_vals = [t.pval.const for t in known_tracers]
    1402        outs1_res = bind(backend.jit_op, *known_vals, program=program1)
    1403        outs1, res = split_list(outs1_res, len(program1.outs) - num_res)
    1404        res_tracers = [trace.instantiate_const(full_raise(trace, x)) for x in res]
    1405        outs2 = [PartialRunTraceTensor(trace, make_unknown_pval(v.symval), None) for v in program2.outs]
    1406        instruction = InstructionDraft(
    1407            self,
    1408            res_tracers + unknown_tracers,
    1409            dict(program=program2),
    1410            [v.symval for v in program2.outs],
    1411            list_map(weakref.ref, outs2),
    1412        )
    1413        for t in outs2:
    1414            t.draft = instruction
    1415
    1416        return merge_lists(out_unknowns, outs1, outs2)
    1417
    1418    def partial_run_instruction(self, unks_in, instruction) -> Tuple[Instruction, Instruction, List[bool], List[Var]]:
    1419        program = instruction.params["program"]
    1420        program1, program2, out_unknowns, num_res = partial_run_program(program, unks_in)
    1421        ins1, ins2 = partition_list(unks_in, instruction.inputs)
    1422        out_binders1, out_binders2 = partition_list(out_unknowns, instruction.out_binders)
    1423        res = [Var(v.symval) for v in program2.in_binders[:num_res]]
    1424        instruction1 = Instruction(self, ins1, dict(program=program1), out_binders1 + res)
    1425        instruction2 = Instruction(self, res + ins2, dict(program=program2), out_binders2)
    1426        return instruction1, instruction2, out_unknowns, res
def meta_impl(self, *args, program: Program, **_): View Source

    1340    def meta_impl(self, *args, program: Program, **_):
    1341        hashed_program = Hashed(program)
    1342        num_consts = program.num_consts
    1343        consts, args = args[:num_consts], args[num_consts:]
    1344        hashed_consts = tuple(map(Hashed, consts))
    1345        jit_output = backend.jit_program(hashed_program, hashed_consts)
    1346        ret = jit_output(*consts, *args)
    1347        return ret
def reorg_args(self, args, params): View Source

    1349    def reorg_args(self, args, params):
    1350        return args, params
def typecheck(self, *in_types, program: Program): View Source

    1352    def typecheck(self, *in_types, program: Program):
    1353        program_type = typecheck_program(program)
    1354        if not all(t1 == t2 for t1, t2 in zip(program_type.in_types, in_types)):
    1355            ret = "Type mismatch program.in_types vs in_types:\n"
    1356            for i, j in zip(program_type.in_types, in_types):
    1357                ret += f"{i}, {j}, {i == j}"
    1358            raise TypeError(ret)
    1359        return program_type.out_types
def vmap(self, dim_size, vals_in, dims_in, program: Program): View Source

    1361    def vmap(self, dim_size, vals_in, dims_in, program: Program):
    1362        program, consts = vmap_program(program, dim_size, tuple(dims_in))
    1363        outs = self(*consts, *vals_in, program=program)
    1364        if not isinstance(outs, tuple):
    1365            outs = (outs,)
    1366        return outs, [0] * len(outs)
def jvp(self, primals, tangents, *, program): View Source

    1368    def jvp(self, primals, tangents, *, program):
    1369        new_program, new_consts = jvp_program(program)
    1370        outs = bind(
    1371            self,
    1372            *new_consts,
    1373            *primals,
    1374            *tangents,
    1375            program=new_program,
    1376        )
    1377        n = len(outs) // 2
    1378        primals_out, tangents_out = outs[:n], outs[n:]
    1379        return primals_out, tangents_out
def T(self, cotangents, *invals, program): View Source

    1381    def T(self, cotangents, *invals, program):
    1382        undef_primals = [isinstance(x, UndefinedPrimal) for x in invals]
    1383        transposed_program, new_consts = transpose_program(program, tuple(undef_primals))
    1384
    1385        residuals, _ = partition_list(undef_primals, invals)
    1386        outs = bind(
    1387            self,
    1388            *new_consts,
    1389            *residuals,
    1390            *cotangents,
    1391            program=transposed_program,
    1392        )
    1393        outs = iter(outs)
    1394
    1395        return [next(outs) if undef else None for undef in undef_primals]
def partial_run(self, trace, tracers, *, program): View Source

    1397    def partial_run(self, trace, tracers, *, program):
    1398        in_unknowns = [not t.pval.is_known for t in tracers]
    1399        program1, program2, out_unknowns, num_res = partial_run_program(program, in_unknowns)
    1400        known_tracers, unknown_tracers = partition_list(in_unknowns, tracers)
    1401        known_vals = [t.pval.const for t in known_tracers]
    1402        outs1_res = bind(backend.jit_op, *known_vals, program=program1)
    1403        outs1, res = split_list(outs1_res, len(program1.outs) - num_res)
    1404        res_tracers = [trace.instantiate_const(full_raise(trace, x)) for x in res]
    1405        outs2 = [PartialRunTraceTensor(trace, make_unknown_pval(v.symval), None) for v in program2.outs]
    1406        instruction = InstructionDraft(
    1407            self,
    1408            res_tracers + unknown_tracers,
    1409            dict(program=program2),
    1410            [v.symval for v in program2.outs],
    1411            list_map(weakref.ref, outs2),
    1412        )
    1413        for t in outs2:
    1414            t.draft = instruction
    1415
    1416        return merge_lists(out_unknowns, outs1, outs2)
def partial_run_instruction( self, unks_in, instruction) -> Tuple[Instruction, Instruction, List[bool], List[Var]]: View Source

    1418    def partial_run_instruction(self, unks_in, instruction) -> Tuple[Instruction, Instruction, List[bool], List[Var]]:
    1419        program = instruction.params["program"]
    1420        program1, program2, out_unknowns, num_res = partial_run_program(program, unks_in)
    1421        ins1, ins2 = partition_list(unks_in, instruction.inputs)
    1422        out_binders1, out_binders2 = partition_list(out_unknowns, instruction.out_binders)
    1423        res = [Var(v.symval) for v in program2.in_binders[:num_res]]
    1424        instruction1 = Instruction(self, ins1, dict(program=program1), out_binders1 + res)
    1425        instruction2 = Instruction(self, res + ins2, dict(program=program2), out_binders2)
    1426        return instruction1, instruction2, out_unknowns, res

##### Inherited Members

Operator
  Operator
  name
  variadic_inputs
  nary_outputs
  args_fixer
  reorg_args_nary

class MainTrace(typing.NamedTuple): View Source

    1434class MainTrace(NamedTuple):
    1435    level: int
    1436    trace_type: Type["Trace"]
    1437    global_data: Optional[Any]

MainTrace(level, trace_type, global_data)

MainTrace( level: int, trace_type: Type[ForwardRef('Trace')], global_data: Optional[Any])

Create new instance of MainTrace(level, trace_type, global_data)

level: int

Alias for field number 0

trace_type: Type[Trace]

Alias for field number 1

global_data: Optional[Any]

Alias for field number 2

##### Inherited Members

builtins.tuple
  index
  count

class Trace: View Source

    1440class Trace:
    1441    main: MainTrace
    1442
    1443    def __init__(self, main: MainTrace) -> None:
    1444        self.main = main
    1445
    1446    def pure(self, val):
    1447        raise NotImplementedError
    1448
    1449    def run_op(self, op, tracers, params):
    1450        raise NotImplementedError
Trace(main: MainTrace) View Source

    1443    def __init__(self, main: MainTrace) -> None:
    1444        self.main = main
main: MainTrace
def pure(self, val): View Source

    1446    def pure(self, val):
    1447        raise NotImplementedError
def run_op(self, op, tracers, params): View Source

    1449    def run_op(self, op, tracers, params):
    1450        raise NotImplementedError
class RunTrace(Trace): View Source

    1453class RunTrace(Trace):
    1454    pure = lambda self, x: x
    1455
    1456    def run_op(self, op: Operator, args, params):
    1457        if isinstance(op, MetaOperator):
    1458            args, params = op.reorg_args(args, params)
    1459            args, params = op.args_fixer(*args, **params)
    1460            ret = op.meta_impl(*args, **params)
    1461        else:
    1462            fn = self.get_fn(op, *tuple(SymbolicTensor.like(a) for a in args), **params)
    1463            # with Timing(f"RUN {op}"):ret = jit(
    1464            ret = jit(
    1465                fn,
    1466                static_argnames=("params",),
    1467                name=jit.get_jit_name(args, params, op.name),
    1468            )(*args, **params)
    1469
    1470        return ret
    1471
    1472    @staticmethod
    1473    @lru_cache_verbose()
    1474    def get_fn(op, *symval_args, **params):
    1475        def fn(*args, **params):
    1476            return [op(*args, **params)]
    1477
    1478        return fn
def pure(self, x): View Source

    1454    pure = lambda self, x: x
def run_op(self, op: Operator, args, params): View Source

    1456    def run_op(self, op: Operator, args, params):
    1457        if isinstance(op, MetaOperator):
    1458            args, params = op.reorg_args(args, params)
    1459            args, params = op.args_fixer(*args, **params)
    1460            ret = op.meta_impl(*args, **params)
    1461        else:
    1462            fn = self.get_fn(op, *tuple(SymbolicTensor.like(a) for a in args), **params)
    1463            # with Timing(f"RUN {op}"):ret = jit(
    1464            ret = jit(
    1465                fn,
    1466                static_argnames=("params",),
    1467                name=jit.get_jit_name(args, params, op.name),
    1468            )(*args, **params)
    1469
    1470        return ret
def get_fn(*args, **kwargs) -> Any: View Source

    172        def decorated_function(*args, **kwargs) -> Any:
    173            result = wrapper(*args, **kwargs)
    174            cache_info = wrapper.cache_info()
    175
    176            dblog(
    177                f"{fn.__name__}.{cache_info} {args.__hash__()}",
    178                enable=backend.LOG_LRU,
    179            )
    180            tb = "".join(traceback.format_list(traceback.extract_stack())[tb_start:tb_end]).replace("\n    ", ":\t") + "-" * 20 + "\n"
    181            dblog(f"{tb}", enable=backend.LOG_LRU)
    182
    183            return result

##### Inherited Members

Trace
  Trace
  main

class SymbolicRunTrace(Trace): View Source

    1481class SymbolicRunTrace(Trace):
    1482    # pure = lambda self, x: x
    1483    def pure(self, val: Any) -> SymbolicTensor:
    1484        return val.symval
    1485
    1486    def run_op(self, op, tracers, params):
    1487        symvals_in = tree_map(lambda x: x.symval, tracers)
    1488        symvals_out = op.typecheck(*symvals_in, **params)
    1489        return symvals_out
def pure(self, val: Any) -> SymbolicTensor: View Source

    1483    def pure(self, val: Any) -> SymbolicTensor:
    1484        return val.symval
def run_op(self, op, tracers, params): View Source

    1486    def run_op(self, op, tracers, params):
    1487        symvals_in = tree_map(lambda x: x.symval, tracers)
    1488        symvals_out = op.typecheck(*symvals_in, **params)
    1489        return symvals_out

##### Inherited Members

Trace
  Trace
  main

class TraceTensor(Tensor): View Source

    1492class TraceTensor(Tensor):
    1493    PYTHON_TYPES = {
    1494        bool,
    1495        int,
    1496        float,
    1497    }
    1498    _trace: "Trace"
    1499
    1500    def __init__(self, *args, **kwargs):
    1501        raise NotImplementedError
    1502
    1503    symval = property(lambda self: get_symval(self.val))
    1504    dtype = property(lambda self: self.symval.dtype)
    1505    shape = property(lambda self: self.symval.shape)
    1506    device = property(lambda self: self.symval.device)
    1507
    1508    @property
    1509    def val(self):
    1510        raise NotImplementedError
    1511
    1512    def __str__(self):
    1513        return repr(self)
    1514
    1515    def full_lower(self):
    1516        return self
    1517
    1518    @property
    1519    def ndim(self):
    1520        return len(self.shape)
    1521
    1522    def __repr__(self):
    1523        return f"{self.__class__.__name__}({repr(self.symval)})"
TraceTensor(*args, **kwargs) View Source

    1500    def __init__(self, *args, **kwargs):
    1501        raise NotImplementedError
PYTHON_TYPES = {<class 'int'>, <class 'float'>, <class 'bool'>}
symval View Source

    1503    symval = property(lambda self: get_symval(self.val))
dtype View Source

    1504    dtype = property(lambda self: self.symval.dtype)
shape View Source

    1505    shape = property(lambda self: self.symval.shape)
device View Source

    1506    device = property(lambda self: self.symval.device)
val View Source

    1508    @property
    1509    def val(self):
    1510        raise NotImplementedError
def full_lower(self): View Source

    1515    def full_lower(self):
    1516        return self
ndim View Source

    1518    @property
    1519    def ndim(self):
    1520        return len(self.shape)

##### Inherited Members

Tensor
  buf
  default_dtype
  is_int
  is_float
  is_unsigned
  to_bool
  short
  int
  long
  half
  float
  str_short
  size
  numpy
  numel
  element_size
  nbytes

class VMapTraceTensor(TraceTensor): View Source

    1526class VMapTraceTensor(TraceTensor):
    1527    def __init__(self, trace, val, vmap_dim):
    1528        self._trace = trace
    1529        self._val = val
    1530        self.vmap_dim = vmap_dim
    1531
    1532    @property
    1533    def val(self):
    1534        return self._val
    1535
    1536    @property
    1537    def symval(self):
    1538        symval = get_symval(self.val)
    1539        if self.vmap_dim is None:
    1540            return symval
    1541        else:
    1542            shape = list(symval.shape)
    1543            del shape[self.vmap_dim]
    1544            return symval.like(shape=tuple(shape))
    1545
    1546    def full_lower(self):
    1547        if self.vmap_dim is None:
    1548            return full_lower(self.val)
    1549        else:
    1550            return self
VMapTraceTensor(trace, val, vmap_dim) View Source

    1527    def __init__(self, trace, val, vmap_dim):
    1528        self._trace = trace
    1529        self._val = val
    1530        self.vmap_dim = vmap_dim
vmap_dim
val View Source

    1532    @property
    1533    def val(self):
    1534        return self._val
symval View Source

    1536    @property
    1537    def symval(self):
    1538        symval = get_symval(self.val)
    1539        if self.vmap_dim is None:
    1540            return symval
    1541        else:
    1542            shape = list(symval.shape)
    1543            del shape[self.vmap_dim]
    1544            return symval.like(shape=tuple(shape))
def full_lower(self): View Source

    1546    def full_lower(self):
    1547        if self.vmap_dim is None:
    1548            return full_lower(self.val)
    1549        else:
    1550            return self

##### Inherited Members

TraceTensor
  PYTHON_TYPES
  dtype
  shape
  device
  ndim
Tensor
  buf
  default_dtype
  is_int
  is_float
  is_unsigned
  to_bool
  short
  int
  long
  half
  float
  str_short
  size
  numpy
  numel
  element_size
  nbytes

class VMapTrace(Trace): View Source

    1553class VMapTrace(Trace):
    1554    pure = lambda self, val: VMapTraceTensor(self, val, None)
    1555
    1556    @property
    1557    def dim_size(self):
    1558        return self.main.global_data
    1559
    1560    def run_op(self, op, tracers, params):
    1561        vals_in, bdims_in = unzip2((t.val, t.vmap_dim) for t in tracers)
    1562        val_outs, bdim_outs = op.vmap(self.dim_size, vals_in, bdims_in, **params)
    1563        return [VMapTraceTensor(self, x, bd) for x, bd in list_zip(val_outs, bdim_outs)]
    1564
    1565    @staticmethod
    1566    def move_vmap_dim(x, dim_size, src: int, dst: int):
    1567        if src is None:  # unsqueeze and expand
    1568            target_shape = list(x.shape)
    1569            target_shape.insert(dst, dim_size)
    1570            unsqueeze_shape = [1 if d == dst else target_shape[d] for d in range(len(target_shape))]
    1571            x = x.reshape(tuple(unsqueeze_shape))
    1572            x = x.expand(tuple(target_shape))
    1573            return x
    1574        elif src == dst:
    1575            return x
    1576        else:
    1577            perm = [i for i in range(len(x.shape)) if i != src]
    1578            perm.insert(dst, src)
    1579            return x.permute(tuple(perm))
def pure(self, val): View Source

    1554    pure = lambda self, val: VMapTraceTensor(self, val, None)
dim_size View Source

    1556    @property
    1557    def dim_size(self):
    1558        return self.main.global_data
def run_op(self, op, tracers, params): View Source

    1560    def run_op(self, op, tracers, params):
    1561        vals_in, bdims_in = unzip2((t.val, t.vmap_dim) for t in tracers)
    1562        val_outs, bdim_outs = op.vmap(self.dim_size, vals_in, bdims_in, **params)
    1563        return [VMapTraceTensor(self, x, bd) for x, bd in list_zip(val_outs, bdim_outs)]
@staticmethod
def move_vmap_dim(x, dim_size, src: int, dst: int): View Source

    1565    @staticmethod
    1566    def move_vmap_dim(x, dim_size, src: int, dst: int):
    1567        if src is None:  # unsqueeze and expand
    1568            target_shape = list(x.shape)
    1569            target_shape.insert(dst, dim_size)
    1570            unsqueeze_shape = [1 if d == dst else target_shape[d] for d in range(len(target_shape))]
    1571            x = x.reshape(tuple(unsqueeze_shape))
    1572            x = x.expand(tuple(target_shape))
    1573            return x
    1574        elif src == dst:
    1575            return x
    1576        else:
    1577            perm = [i for i in range(len(x.shape)) if i != src]
    1578            perm.insert(dst, src)
    1579            return x.permute(tuple(perm))

##### Inherited Members

Trace
  Trace
  main

class JVPTraceTensor(TraceTensor): View Source

    1582class JVPTraceTensor(TraceTensor):
    1583    def __init__(self, trace, primal, tangent):
    1584        self._trace = trace
    1585        self.primal = primal
    1586        self.tangent = tangent
    1587
    1588    @property
    1589    def symval(self):
    1590        return get_symval(self.primal)
    1591
    1592    @property
    1593    def val(self):
    1594        return self.primal
    1595
    1596    @property
    1597    def dtype(self):
    1598        return self.primal.dtype
JVPTraceTensor(trace, primal, tangent) View Source

    1583    def __init__(self, trace, primal, tangent):
    1584        self._trace = trace
    1585        self.primal = primal
    1586        self.tangent = tangent
primal
tangent
symval View Source

    1588    @property
    1589    def symval(self):
    1590        return get_symval(self.primal)
val View Source

    1592    @property
    1593    def val(self):
    1594        return self.primal
dtype View Source

    1596    @property
    1597    def dtype(self):
    1598        return self.primal.dtype

##### Inherited Members

TraceTensor
  PYTHON_TYPES
  shape
  device
  full_lower
  ndim
Tensor
  buf
  default_dtype
  is_int
  is_float
  is_unsigned
  to_bool
  short
  int
  long
  half
  float
  str_short
  size
  numpy
  numel
  element_size
  nbytes

class JVPTrace(Trace): View Source

    1601class JVPTrace(Trace):
    1602    def pure(self, val):
    1603        if isinstance(val, PartialRunTrace):
    1604            val = val.pval.const
    1605        return JVPTraceTensor(self, val, backend.zeros_like(val))
    1606
    1607    def run_op(self, op, tracers, params):
    1608        primals_in, tangents_in = unzip2((t.primal, t.tangent) for t in tracers)
    1609        primals_out, tangents_out = op.jvp(primals_in, tangents_in, **params)
    1610        return [JVPTraceTensor(self, x, t) for x, t in list_zip(primals_out, tangents_out)]
def pure(self, val): View Source

    1602    def pure(self, val):
    1603        if isinstance(val, PartialRunTrace):
    1604            val = val.pval.const
    1605        return JVPTraceTensor(self, val, backend.zeros_like(val))
def run_op(self, op, tracers, params): View Source

    1607    def run_op(self, op, tracers, params):
    1608        primals_in, tangents_in = unzip2((t.primal, t.tangent) for t in tracers)
    1609        primals_out, tangents_out = op.jvp(primals_in, tangents_in, **params)
    1610        return [JVPTraceTensor(self, x, t) for x, t in list_zip(primals_out, tangents_out)]

##### Inherited Members

Trace
  Trace
  main

class ProgramTraceTensor(TraceTensor): View Source

    1613class ProgramTraceTensor(TraceTensor):
    1614    __slots__ = ["symval"]
    1615    symval: SymbolicTensor
    1616
    1617    def __init__(self, trace, symval):
    1618        self._trace = trace
    1619        self.symval = symval
ProgramTraceTensor(trace, symval) View Source

    1617    def __init__(self, trace, symval):
    1618        self._trace = trace
    1619        self.symval = symval
symval: SymbolicTensor

##### Inherited Members

TraceTensor
  PYTHON_TYPES
  dtype
  shape
  device
  val
  full_lower
  ndim
Tensor
  buf
  default_dtype
  is_int
  is_float
  is_unsigned
  to_bool
  short
  int
  long
  half
  float
  str_short
  size
  numpy
  numel
  element_size
  nbytes

class ProgramTrace(Trace): View Source

    1622class ProgramTrace(Trace):
    1623    @property
    1624    def builder(self):
    1625        return self.main.global_data
    1626
    1627    def new_arg(self, symval) -> ProgramTraceTensor:
    1628        symval = SymbolicTensor.like(symval)
    1629        tracer = self.builder.new_tracer(self, symval)
    1630        self.builder.tracer_to_var[id(tracer)] = Var(symval)
    1631
    1632        return tracer
    1633
    1634    def pure(self, val: Any) -> ProgramTraceTensor:
    1635        # get_or_make_const_tracer
    1636        tracer = self.builder.const_tracers.get(id(val))
    1637        if tracer is None:
    1638            tracer = self.builder.new_tracer(self, get_symval(val))
    1639            self.builder.add_const(tracer, val)
    1640        # print(self.builder.const_tracers)
    1641        return tracer
    1642
    1643    def run_op(self, op, tracers, params):
    1644        symvals_in = tree_map(lambda x: x.symval, tracers)
    1645        symvals_out = op.typecheck(*symvals_in, **params)
    1646
    1647        out_tracers = [self.builder.new_tracer(self, a) for a in symvals_out]
    1648        inputs = [self.builder.getvar(t) for t in tracers]
    1649        outvars = [self.builder.add_var(t) for t in out_tracers]
    1650
    1651        self.builder.add_instruction(Instruction(op, inputs, params, outvars))
    1652        return out_tracers
builder View Source

    1623    @property
    1624    def builder(self):
    1625        return self.main.global_data
def new_arg(self, symval) -> ProgramTraceTensor: View Source

    1627    def new_arg(self, symval) -> ProgramTraceTensor:
    1628        symval = SymbolicTensor.like(symval)
    1629        tracer = self.builder.new_tracer(self, symval)
    1630        self.builder.tracer_to_var[id(tracer)] = Var(symval)
    1631
    1632        return tracer
def pure(self, val: Any) -> ProgramTraceTensor: View Source

    1634    def pure(self, val: Any) -> ProgramTraceTensor:
    1635        # get_or_make_const_tracer
    1636        tracer = self.builder.const_tracers.get(id(val))
    1637        if tracer is None:
    1638            tracer = self.builder.new_tracer(self, get_symval(val))
    1639            self.builder.add_const(tracer, val)
    1640        # print(self.builder.const_tracers)
    1641        return tracer
def run_op(self, op, tracers, params): View Source

    1643    def run_op(self, op, tracers, params):
    1644        symvals_in = tree_map(lambda x: x.symval, tracers)
    1645        symvals_out = op.typecheck(*symvals_in, **params)
    1646
    1647        out_tracers = [self.builder.new_tracer(self, a) for a in symvals_out]
    1648        inputs = [self.builder.getvar(t) for t in tracers]
    1649        outvars = [self.builder.add_var(t) for t in out_tracers]
    1650
    1651        self.builder.add_instruction(Instruction(op, inputs, params, outvars))
    1652        return out_tracers

##### Inherited Members

Trace
  Trace
  main

class ProgramBuilder: View Source

    1655class ProgramBuilder:
    1656    instructions: List[Instruction]
    1657    tracer_to_var: Dict[int, Var]
    1658    const_tracers: Dict[int, TraceTensor]
    1659    constvals: Dict[Var, Any]
    1660    tracers: List[ProgramTraceTensor]
    1661
    1662    def __init__(self):
    1663        self.instructions = []
    1664        self.tracer_to_var = {}
    1665        self.const_tracers = {}
    1666        self.constvals = {}
    1667        self.tracers = []
    1668
    1669    def new_tracer(self, trace: ProgramTrace, symval: SymbolicTensor) -> ProgramTraceTensor:
    1670        tracer = ProgramTraceTensor(trace, symval)
    1671        self.tracers += [tracer]
    1672        return tracer
    1673
    1674    def add_instruction(self, instruction: Instruction) -> None:
    1675        self.instructions += [instruction]
    1676
    1677    def add_var(self, tracer: ProgramTraceTensor) -> Var:
    1678        assert id(tracer) not in self.tracer_to_var
    1679        var = self.tracer_to_var[id(tracer)] = Var(tracer.symval)
    1680        return var
    1681
    1682    def getvar(self, tracer: ProgramTraceTensor) -> Var:
    1683        var = self.tracer_to_var.get(id(tracer))
    1684        assert var is not None
    1685        return var
    1686
    1687    def add_const(self, tracer: ProgramTraceTensor, val: Any) -> Var:
    1688        var = self.add_var(tracer)
    1689        self.const_tracers[id(val)] = tracer
    1690        self.constvals[var] = val
    1691        return var
    1692
    1693    def build(self, in_tracers: Any, out_tracers: Any, static_args, name) -> Tuple[Program, List[Any]]:
    1694        constvars, constvals = unzip2(self.constvals.items())
    1695        t2v = lambda t: self.tracer_to_var[id(t)]
    1696        in_binders = constvars + [t2v(t) for t in in_tracers]
    1697        out_vars = [t2v(t) for t in out_tracers]
    1698        program = Program(
    1699            in_binders,
    1700            self.instructions,
    1701            out_vars,
    1702            len(constvals),
    1703            static_args,
    1704            name,
    1705        )
    1706        typecheck_program(program)
    1707        program, constvals = self._inline_literals(program, constvals)
    1708        typecheck_program(program)
    1709        # dblog(program, enable=backend.LOG_PROGRAM)
    1710        return program, constvals
    1711
    1712    def _inline_literals(self, program: Program, consts: List[Any]) -> Tuple[Program, List[Any]]:
    1713        const_binders, other_binders = split_list(program.in_binders, len(consts))
    1714        scalars = [type(x) in TraceTensor.PYTHON_TYPES and not get_symval(x).shape for x in consts]
    1715        new_const_binders, lit_binders = partition_list(scalars, const_binders)
    1716        new_consts, lit_vals = partition_list(scalars, consts)
    1717        literals = dict(list_zip(lit_binders, list_map(Lit, lit_vals)))
    1718        new_outs = [literals.get(x, x) for x in program.outs]
    1719        new_instructions = [
    1720            Instruction(
    1721                instruction.op,
    1722                [literals.get(x, x) for x in instruction.inputs],
    1723                instruction.params,
    1724                instruction.out_binders,
    1725            )
    1726            for instruction in program.instructions
    1727        ]
    1728        new_program = Program(
    1729            new_const_binders + other_binders,
    1730            new_instructions,
    1731            new_outs,
    1732            len(new_consts),
    1733            program.static_args,
    1734            program.name,
    1735        )
    1736        return new_program, tuple(new_consts)
    1737
    1738    def get_current_scope_info(self):
    1739        current_frame = inspect.currentframe()
    1740        current_function_name = current_frame.f_code.co_name
    1741        current_module_name = inspect.getmodulename(current_frame.f_code.co_filename)
    1742        current_class_name = None
    1743        for frame_info in inspect.getouterframes(current_frame):
    1744            print(frame_info)
    1745            frame_locals = frame_info.frame.f_locals
    1746            print(frame_locals)
    1747            if "self" in frame_locals:
    1748                current_class_name = frame_locals["self"].__class__.__name__
    1749                break
    1750        return {
    1751            "Function": current_function_name,
    1752            "Module": current_module_name,
    1753            "Class": current_class_name,
    1754        }
instructions: List[Instruction]
tracer_to_var: Dict[int, Var]
const_tracers: Dict[int, TraceTensor]
constvals: Dict[Var, Any]
tracers: List[ProgramTraceTensor]
def new_tracer( self, trace: ProgramTrace, symval: SymbolicTensor) -> ProgramTraceTensor: View Source

    1669    def new_tracer(self, trace: ProgramTrace, symval: SymbolicTensor) -> ProgramTraceTensor:
    1670        tracer = ProgramTraceTensor(trace, symval)
    1671        self.tracers += [tracer]
    1672        return tracer
def add_instruction(self, instruction: Instruction) -> None: View Source

    1674    def add_instruction(self, instruction: Instruction) -> None:
    1675        self.instructions += [instruction]
def add_var(self, tracer: ProgramTraceTensor) -> Var: View Source

    1677    def add_var(self, tracer: ProgramTraceTensor) -> Var:
    1678        assert id(tracer) not in self.tracer_to_var
    1679        var = self.tracer_to_var[id(tracer)] = Var(tracer.symval)
    1680        return var
def getvar(self, tracer: ProgramTraceTensor) -> Var: View Source

    1682    def getvar(self, tracer: ProgramTraceTensor) -> Var:
    1683        var = self.tracer_to_var.get(id(tracer))
    1684        assert var is not None
    1685        return var
def add_const(self, tracer: ProgramTraceTensor, val: Any) -> Var: View Source

    1687    def add_const(self, tracer: ProgramTraceTensor, val: Any) -> Var:
    1688        var = self.add_var(tracer)
    1689        self.const_tracers[id(val)] = tracer
    1690        self.constvals[var] = val
    1691        return var
def build( self, in_tracers: Any, out_tracers: Any, static_args, name) -> Tuple[Program, List[Any]]: View Source

    1693    def build(self, in_tracers: Any, out_tracers: Any, static_args, name) -> Tuple[Program, List[Any]]:
    1694        constvars, constvals = unzip2(self.constvals.items())
    1695        t2v = lambda t: self.tracer_to_var[id(t)]
    1696        in_binders = constvars + [t2v(t) for t in in_tracers]
    1697        out_vars = [t2v(t) for t in out_tracers]
    1698        program = Program(
    1699            in_binders,
    1700            self.instructions,
    1701            out_vars,
    1702            len(constvals),
    1703            static_args,
    1704            name,
    1705        )
    1706        typecheck_program(program)
    1707        program, constvals = self._inline_literals(program, constvals)
    1708        typecheck_program(program)
    1709        # dblog(program, enable=backend.LOG_PROGRAM)
    1710        return program, constvals
def get_current_scope_info(self): View Source

    1738    def get_current_scope_info(self):
    1739        current_frame = inspect.currentframe()
    1740        current_function_name = current_frame.f_code.co_name
    1741        current_module_name = inspect.getmodulename(current_frame.f_code.co_filename)
    1742        current_class_name = None
    1743        for frame_info in inspect.getouterframes(current_frame):
    1744            print(frame_info)
    1745            frame_locals = frame_info.frame.f_locals
    1746            print(frame_locals)
    1747            if "self" in frame_locals:
    1748                current_class_name = frame_locals["self"].__class__.__name__
    1749                break
    1750        return {
    1751            "Function": current_function_name,
    1752            "Module": current_module_name,
    1753            "Class": current_class_name,
    1754        }
class UndefinedPrimal(typing.NamedTuple): View Source

    1757class UndefinedPrimal(NamedTuple):
    1758    symval: SymbolicTensor
    1759
    1760    @property
    1761    def shape(self):
    1762        return self.symval.shape
    1763
    1764    @property
    1765    def dtype(self):
    1766        return self.symval.dtype
    1767
    1768    @property
    1769    def device(self):
    1770        return self.symval.device
    1771
    1772    @property
    1773    def ndim(self):
    1774        return self.symval.ndim
    1775
    1776    def __repr__(self):
    1777        return f"<UndefinedPrimal: symval={self.symval}>"
    1778
    1779    str_short = __repr__

UndefinedPrimal(symval,)

UndefinedPrimal(symval: SymbolicTensor)

Create new instance of UndefinedPrimal(symval,)

symval: SymbolicTensor

Alias for field number 0

shape View Source

    1760    @property
    1761    def shape(self):
    1762        return self.symval.shape
dtype View Source

    1764    @property
    1765    def dtype(self):
    1766        return self.symval.dtype
device View Source

    1768    @property
    1769    def device(self):
    1770        return self.symval.device
ndim View Source

    1772    @property
    1773    def ndim(self):
    1774        return self.symval.ndim
def str_short(self): View Source

    1776    def __repr__(self):
    1777        return f"<UndefinedPrimal: symval={self.symval}>"

Return repr(self).

##### Inherited Members

builtins.tuple
  index
  count

class PartialValue(typing.NamedTuple): View Source

    1782class PartialValue(NamedTuple):
    1783    symval: SymbolicTensor
    1784    const: Optional[Any]
    1785
    1786    is_known = property(lambda self: self.const is not None)
    1787    is_unknown = property(lambda self: self.const is None)

PartialValue(symval, const)

PartialValue(symval: SymbolicTensor, const: Optional[Any])

Create new instance of PartialValue(symval, const)

symval: SymbolicTensor

Alias for field number 0

const: Optional[Any]

Alias for field number 1

is_known View Source

    1786    is_known = property(lambda self: self.const is not None)
is_unknown View Source

    1787    is_unknown = property(lambda self: self.const is None)

##### Inherited Members

builtins.tuple
  index
  count

class LambdaBindingDraft(typing.NamedTuple): View Source

    1790class LambdaBindingDraft(NamedTuple):
    1791    pass

LambdaBindingDraft()

LambdaBindingDraft()

Create new instance of LambdaBindingDraft()

##### Inherited Members

builtins.tuple
  index
  count

class ConstDraft(typing.NamedTuple): View Source

    1794class ConstDraft(NamedTuple):
    1795    val: Any

ConstDraft(val,)

ConstDraft(val: Any)

Create new instance of ConstDraft(val,)

val: Any

Alias for field number 0

##### Inherited Members

builtins.tuple
  index
  count

class InstructionDraft(typing.NamedTuple): View Source

    1798class InstructionDraft(NamedTuple):
    1799    prim: Operator
    1800    tracers_in: List["PartialRunTraceTensor"]
    1801    params: Dict[str, Any]
    1802    symvals_out: List[SymbolicTensor]
    1803    tracer_refs_out: List[weakref.ReferenceType["PartialRunTraceTensor"]]

InstructionDraft(prim, tracers_in, params, symvals_out, tracer_refs_out)

InstructionDraft( prim: Operator, tracers_in: List[ForwardRef('PartialRunTraceTensor')], params: Dict[str, Any], symvals_out: List[SymbolicTensor], tracer_refs_out: List[weakref.ReferenceType['PartialRunTraceTensor']])

Create new instance of InstructionDraft(prim, tracers_in, params, symvals_out, tracer_refs_out)

prim: Operator

Alias for field number 0

tracers_in: List[PartialRunTraceTensor]

Alias for field number 1

params: Dict[str, Any]

Alias for field number 2

symvals_out: List[SymbolicTensor]

Alias for field number 3

tracer_refs_out: List[weakref.ReferenceType[PartialRunTraceTensor]]

Alias for field number 4

##### Inherited Members

builtins.tuple
  index
  count

ProgramDraft = typing.Union[LambdaBindingDraft, ConstDraft, InstructionDraft]
class PartialRunTraceTensor(TraceTensor): View Source

    1809class PartialRunTraceTensor(TraceTensor):
    1810    def __init__(self, trace, pval, draft):
    1811        self._trace = trace
    1812        self.pval = pval
    1813        self.draft = draft
    1814
    1815    symval = property(lambda self: self.pval.symval)
    1816    val = property(lambda self: self.pval.const)
    1817
    1818    def full_lower(self):
    1819        if self.pval.is_known:
    1820            return full_lower(self.pval.const)
    1821        return self
PartialRunTraceTensor(trace, pval, draft) View Source

    1810    def __init__(self, trace, pval, draft):
    1811        self._trace = trace
    1812        self.pval = pval
    1813        self.draft = draft
pval
draft
symval View Source

    1815    symval = property(lambda self: self.pval.symval)
val View Source

    1816    val = property(lambda self: self.pval.const)
def full_lower(self): View Source

    1818    def full_lower(self):
    1819        if self.pval.is_known:
    1820            return full_lower(self.pval.const)
    1821        return self

##### Inherited Members

TraceTensor
  PYTHON_TYPES
  dtype
  shape
  device
  ndim
Tensor
  buf
  default_dtype
  is_int
  is_float
  is_unsigned
  to_bool
  short
  int
  long
  half
  float
  str_short
  size
  numpy
  numel
  element_size
  nbytes

class PartialRunTrace(Trace): View Source

    1824class PartialRunTrace(Trace):
    1825    def new_arg(self, pval: PartialValue) -> Any:
    1826        return PartialRunTraceTensor(self, pval, LambdaBindingDraft())
    1827
    1828    def pure(self, val: Any) -> PartialRunTraceTensor:
    1829        return PartialRunTraceTensor(self, make_known_pval(val), None)
    1830
    1831    def instantiate_const(self, tracer: PartialRunTraceTensor) -> PartialRunTraceTensor:
    1832        if tracer.pval.is_unknown:
    1833            return tracer
    1834        else:
    1835            pval = make_unknown_pval(SymbolicTensor.like(tracer.symval))
    1836            return PartialRunTraceTensor(self, pval, ConstDraft(tracer.pval.const))
    1837
    1838    def run_op(self, op, tracers, params):
    1839        is_knowns = tuple(t.pval.is_known for t in tracers)
    1840
    1841        if all(is_knowns):
    1842            return bind(op, *list_map(full_lower, tracers), **params)
    1843        return op.partial_run(self, tracers, **params)
def new_arg(self, pval: PartialValue) -> Any: View Source

    1825    def new_arg(self, pval: PartialValue) -> Any:
    1826        return PartialRunTraceTensor(self, pval, LambdaBindingDraft())
def pure(self, val: Any) -> PartialRunTraceTensor: View Source

    1828    def pure(self, val: Any) -> PartialRunTraceTensor:
    1829        return PartialRunTraceTensor(self, make_known_pval(val), None)
def instantiate_const( self, tracer: PartialRunTraceTensor) -> PartialRunTraceTensor: View Source

    1831    def instantiate_const(self, tracer: PartialRunTraceTensor) -> PartialRunTraceTensor:
    1832        if tracer.pval.is_unknown:
    1833            return tracer
    1834        else:
    1835            pval = make_unknown_pval(SymbolicTensor.like(tracer.symval))
    1836            return PartialRunTraceTensor(self, pval, ConstDraft(tracer.pval.const))
def run_op(self, op, tracers, params): View Source

    1838    def run_op(self, op, tracers, params):
    1839        is_knowns = tuple(t.pval.is_known for t in tracers)
    1840
    1841        if all(is_knowns):
    1842            return bind(op, *list_map(full_lower, tracers), **params)
    1843        return op.partial_run(self, tracers, **params)

##### Inherited Members

Trace
  Trace
  main

trace_stack: List[MainTrace] = [MainTrace(level=0, trace_type=<class 'RunTrace'>, global_data=None)]
stashed_trace: Optional[MainTrace] = None
class UndefBackend: View Source

    1851class UndefBackend:
    1852    def __getattr__(self, attr):
    1853        raise NotImplementedError("Backend not init yet with slope.core.set_backend(backend)")
backend = <slope.backends.onnxruntime.ONNXRuntimeBackend object>
def set_backend(name, where='slope.backends'): View Source

    1859def set_backend(name, where="slope.backends"):
    1860    global backend
    1861    backend = importlib.import_module(f"{where}.{name}").backend
    1862    import slope.nn as nn
    1863    # backend.register_node(nn.Module, nn.Module.flatten, nn.Module.unflatten, "Module")
    1864
    1865    dblog(f"slope backend is {backend}", enable=backend.LOG_INIT)
def stack_str(): View Source

    1868def stack_str():
    1869    ret = ""
    1870    for trace in trace_stack:
    1871        ret += f"{trace.level}: {trace.trace_type.__name__}\t{trace.global_data=}\n"
    1872    return ret
def make_known_pval(val: Any): View Source

    1875def make_known_pval(val: Any):
    1876    return PartialValue(get_symval(val), val)
def make_unknown_pval(symval: SymbolicTensor): View Source

    1879def make_unknown_pval(symval: SymbolicTensor):
    1880    return PartialValue(symval, None)
def get_symval(x): View Source

    1883def get_symval(x):
    1884    if isinstance(x, TraceTensor):
    1885        return x.symval
    1886    elif type(x) in TraceTensor.PYTHON_TYPES:
    1887        return backend.tensor(x)
    1888    elif isinstance(x, Tensor):
    1889        return x
    1890    elif isinstance(x, SymbolicTensor):
    1891        return x
    1892    else:
    1893        raise TypeError(type(x))
def tree_flatten(x: Any) -> Any: View Source

    1896def tree_flatten(x: Any) -> Any:
    1897    def _tree_flatten(x_: Any) -> Tuple[Iterable, Union[TreeDef, Leaf]]:
    1898        node_type = None
    1899        for k in backend.node_types.keys():
    1900            if isinstance(x_, k):
    1901                node_type = backend.node_types[k]
    1902
    1903        if node_type is not None:
    1904            node_metadata, children = node_type.flatten(x_)
    1905            children_flat, child_trees = unzip2(list_map(_tree_flatten, children))
    1906            children_iter = itertools.chain.from_iterable(children_flat)
    1907            treedef = TreeDef(node_type, node_metadata, tuple(child_trees))
    1908            return children_iter, treedef
    1909        else:
    1910            return (x_,), Leaf(x_)
    1911
    1912    children_iter, treedef = _tree_flatten(x)
    1913    return tuple(children_iter), treedef
def tree_unflatten(treedef: TreeDef, xs: Tuple[Any]) -> Any: View Source

    1916def tree_unflatten(treedef: TreeDef, xs: Tuple[Any]) -> Any:
    1917    def _tree_unflatten(treedef_: TreeDef, xs_: Iterator) -> Any:
    1918        if isinstance(treedef_, Leaf):
    1919            dblog(f"    tree leaf found: {xs_}\n", enable=backend.LOG_TREE)
    1920            return next(xs_)
    1921        else:
    1922            dblog(f"    now\n  {treedef_}", enable=backend.LOG_TREE)
    1923            children = (_tree_unflatten(t, xs_) for t in treedef_.child_treedefs)
    1924            dblog(f"{children=}\n", enable=backend.LOG_TREE)
    1925            return treedef_.node_type.unflatten(treedef_.node_metadata, children)
    1926
    1927    dblog(f"unflattening {treedef}", enable=backend.LOG_TREE)
    1928    return _tree_unflatten(treedef, iter(xs))
    1929    # with Timing(f"\nTREE:\n{treedef}"):
    1930    #     ret = _tree_unflatten(treedef, iter(xs))
    1931    # return ret
def tree_transpose( outer_treedef: TreeDef, inner_treedef: TreeDef, tree_to_transpose: Any) -> Any: View Source

    1934def tree_transpose(
    1935    outer_treedef: TreeDef,
    1936    inner_treedef: TreeDef,
    1937    tree_to_transpose: Any,
    1938) -> Any:
    1939    flat, treedef = tree_flatten(tree_to_transpose)
    1940    inner_size = inner_treedef.num_leaves
    1941    outer_size = outer_treedef.num_leaves
    1942    if treedef.num_leaves != (inner_size * outer_size):
    1943        raise TypeError
    1944    iter_flat = iter(flat)
    1945    lol = [[next(iter_flat) for _ in range(inner_size)] for __ in range(outer_size)]
    1946    permuted_lol = zip(*lol)
    1947    subtrees = map(partial(tree_unflatten, outer_treedef), permuted_lol)
    1948    return tree_unflatten(inner_treedef, subtrees)
def flatten_fn(f, in_tree, *, has_aux=False): View Source

    1951def flatten_fn(f, in_tree, *, has_aux=False):
    1952    store = Store()
    1953
    1954    def flat_fn(*args_flat, **params):
    1955        tree_args = tree_unflatten(in_tree, args_flat)
    1956        out = f(*tree_args, **params)
    1957        if has_aux:
    1958            out, aux = out
    1959        out_flat, out_tree = tree_flatten(out)
    1960        store.set_value(out_tree)
    1961        return (out_flat, aux) if has_aux else out_flat
    1962
    1963    return flat_fn, store
def tree_map(f: Callable[..., Any], tree, *rest, out_leaf=False) -> Any: View Source

    1966def tree_map(f: Callable[..., Any], tree, *rest, out_leaf=False) -> Any:
    1967    leaves, treedef = tree_flatten(tree)
    1968    if len(rest) == 0:
    1969        out_tree_flat = tuple(f(leaf) for leaf in leaves)
    1970        out_tree = tree_unflatten(treedef, out_tree_flat)
    1971    else:
    1972        all_leaves = [leaves]
    1973        for t in rest:
    1974            t_leaves, t_treedef = tree_flatten(t)
    1975            assert t_treedef == treedef
    1976            all_leaves += [t_leaves]
    1977
    1978        out_tree_flat = tuple(f(*xs) for xs in zip(*all_leaves))
    1979        out_tree = tree_unflatten(treedef, out_tree_flat)
    1980    ret = out_tree
    1981    if out_leaf:
    1982        ret = (ret, tree_flatten(out_tree_flat[0]))
    1983    return ret
@contextmanager
def new_main_trace(trace_type: Type[Trace], global_data=None): View Source

    1986@contextmanager
    1987def new_main_trace(trace_type: Type["Trace"], global_data=None):
    1988    global trace_stack
    1989    level = len(trace_stack)
    1990    main = MainTrace(level, trace_type, global_data)
    1991    trace_stack += [main]
    1992
    1993    try:
    1994        yield main
    1995    finally:
    1996        trace_stack.pop()
def bind(op, *args, **params): View Source

    1999def bind(op, *args, **params):
    2000    top_trace = find_top_trace(args)
    2001    tracers = tuple([full_raise(top_trace, arg) for arg in args])
    2002    outs = top_trace.run_op(op, tracers, params)
    2003    lowered = tuple([full_lower(out) for out in outs])
    2004    return lowered
def find_top_trace(xs) -> Trace: View Source

    2007def find_top_trace(xs) -> Trace:
    2008    arrs = []
    2009
    2010    def get_arr_from_seq(seq):
    2011        nonlocal arrs
    2012        for x in seq:
    2013            if type(x) in (tuple, list):
    2014                get_arr_from_seq(x)
    2015            elif isinstance(x, TraceTensor):
    2016                arrs += [x]
    2017
    2018    get_arr_from_seq(xs)
    2019    arrs = tuple(arrs)
    2020    top_main = max(
    2021        (x._trace.main for x in arrs),
    2022        default=trace_stack[0],
    2023        key=operator_py.attrgetter("level"),
    2024    )
    2025    if stashed_trace and stashed_trace.level > top_main.level:
    2026        top_main = stashed_trace
    2027    return top_main.trace_type(top_main)
def full_raise(trace: Trace, val: Any) -> TraceTensor: View Source

    2030def full_raise(trace: Trace, val: Any) -> TraceTensor:
    2031    if not isinstance(val, TraceTensor):
    2032        return trace.pure(val)
    2033    level = trace.main.level
    2034    if val._trace.main is trace.main:
    2035        return val
    2036    elif val._trace.main.level < level:
    2037        return trace.pure(val)
    2038    elif val._trace.main.level > level:
    2039        raise Exception(f"Can't lift level {val._trace.main.level} to {level}.")
    2040    else:
    2041        raise Exception(f"Different traces at same level: {val._trace}, {trace}.")
def full_lower(val: Any): View Source

    2044def full_lower(val: Any):
    2045    if isinstance(val, TraceTensor):
    2046        return val.full_lower()
    2047    elif type(val) in (list, tuple):
    2048        return tuple(full_lower(v) for v in val)
    2049    else:
    2050        return val
def typecheck_program(program: Program) -> ProgramType: View Source

    2053def typecheck_program(program: Program) -> ProgramType:
    2054    env: Set[Var] = set()
    2055
    2056    for v in program.in_binders:
    2057        if v in env:
    2058            raise TypeError
    2059        env.add(v)
    2060
    2061    for instruction in program.instructions:
    2062        in_types = [typecheck_atom(env, x) for x in instruction.inputs]
    2063        out_types = instruction.op.typecheck(*in_types, **instruction.params)
    2064        for out_binder, out_type in list_zip(instruction.out_binders, out_types):
    2065            if not out_type == out_binder.symval:
    2066                raise TypeError
    2067        for out_binder in instruction.out_binders:
    2068            if out_binder in env:
    2069                raise TypeError
    2070            env.add(out_binder)
    2071
    2072    in_types = [v.symval for v in program.in_binders]
    2073    out_types = [typecheck_atom(env, x) for x in program.outs]
    2074    return ProgramType(tuple(in_types), tuple(out_types))
def typecheck_atom( env: Set[Var], x: Union[Var, Lit]) -> SymbolicTensor: View Source

    2077def typecheck_atom(env: Set[Var], x: Atom) -> SymbolicTensor:
    2078    if isinstance(x, Var):
    2079        if x not in env:
    2080            raise TypeError("unbound variable")
    2081        return x.symval
    2082    elif isinstance(x, Lit):
    2083        return get_symval(x.val)
    2084    else:
    2085        assert False
def run_program(program: Program, args: List[Any]) -> List[Any]: View Source

    2088def run_program(program: Program, args: List[Any]) -> List[Any]:
    2089    env: Dict[Var, Any] = {}
    2090
    2091    def read(x: Atom) -> Any:
    2092        return env[x] if type(x) is Var else x.val
    2093
    2094    def write(v: Var, val: Any) -> None:
    2095        assert v not in env  # single-assignment
    2096        env[v] = val
    2097
    2098    list_map(write, program.in_binders, args)
    2099    for instruction in program.instructions:
    2100        in_vals = list_map(read, instruction.inputs)
    2101        outs = bind(instruction.op, *in_vals, **instruction.params)
    2102        list_map(write, instruction.out_binders, outs)
    2103    return list_map(read, program.outs)
def program_as_fun(program: Program): View Source

    2106def program_as_fun(program: Program):
    2107    return lambda *args: run_program(program, args)
def vmap_flat(f, in_dim, out_dim, dim_size, *args): View Source

    2110def vmap_flat(f, in_dim, out_dim, dim_size, *args):
    2111    if dim_size is None:
    2112        dims = set([x.shape[d] for x, d in list_zip(args, in_dim) if d is not None])
    2113        assert len(dims) == 1
    2114        (dim_size,) = dims
    2115    with new_main_trace(VMapTrace, dim_size) as main:
    2116        trace = VMapTrace(main)
    2117        tracers_in = [VMapTraceTensor(trace, x, dim) if dim is not None else x for x, dim in list_zip(args, in_dim)]
    2118        outs = f(*tracers_in)
    2119        tracers_out = [full_raise(trace, out) for out in outs]
    2120        vals_out, y_vmap_dims = unzip2((t.val, t.vmap_dim) for t in tracers_out)
    2121    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)]
    2122    return ret
def vmap(f, in_dim=0, out_dim=0, dim_size=None): View Source

    2125def vmap(f, in_dim=0, out_dim=0, dim_size=None):
    2126    def batched_f(*args):
    2127        nonlocal in_dim, out_dim, dim_size
    2128        args_flat, in_tree = tree_flatten(args)
    2129        in_dim = (in_dim,) * len(args) if isinstance(in_dim, int) else in_dim
    2130        out_dim = (out_dim,) * len(args) if isinstance(out_dim, int) else out_dim
    2131        in_dim_flat, in_dim_tree = tree_flatten(in_dim)
    2132        out_dim_flat, out_dim_tree = tree_flatten(out_dim)
    2133        if not (in_tree == in_dim_tree == out_dim_tree):
    2134            raise TypeError(f"\n{in_tree}\n!=\n{in_dim_tree}!=\n{out_dim_tree}")
    2135        f_flat, out_tree_store = flatten_fn(f, in_tree)
    2136        # if len(args_flat) > len(in_dim_flat):
    2137        #     in_dim_flat = (in_dim[0],) * len(args_flat)
    2138        outs_flat = vmap_flat(f_flat, in_dim_flat, out_dim_flat, dim_size, *args_flat)
    2139        return tree_unflatten(out_tree_store(), outs_flat)
    2140
    2141    return batched_f
def jvp_flat(f, primals, tangents, *, has_aux, global_data, **static_args): View Source

    2144def jvp_flat(f, primals, tangents, *, has_aux, global_data, **static_args):
    2145    with new_main_trace(JVPTrace, global_data) as main:
    2146        trace = JVPTrace(main)
    2147        tracers_in = [JVPTraceTensor(trace, x, t) for x, t in list_zip(primals, tangents)]
    2148        jvp_flat_ret = f(*tracers_in, **static_args)
    2149        if has_aux:
    2150            (outs, aux) = jvp_flat_ret
    2151            # aux_ = aux
    2152            aux = tree_map(lambda x: x.primal, aux)
    2153            # aux = tree_map(lambda x: x.full_lower(), aux)
    2154            #
    2155        else:
    2156            outs = jvp_flat_ret
    2157        tracers_out = [full_raise(trace, out) for out in outs]
    2158        primals_out, tangents_out = unzip2((t.primal, t.tangent) for t in tracers_out)
    2159    return ((primals_out, tangents_out), aux) if has_aux else (primals_out, tangents_out)
def jvp( f, primals, tangents, *, has_aux=False, global_data=None, **static_args): View Source

    2162def jvp(f, primals, tangents, *, has_aux=False, global_data=None, **static_args):
    2163    primals_flat, in_tree = tree_flatten(primals)
    2164    tangents_flat, in_tree2 = tree_flatten(tangents)
    2165    for p, t in zip(primals_flat, tangents_flat):
    2166        assert p.shape == t.shape, f"{p.shape=} != {t.shape=}"
    2167        assert p.dtype == t.dtype, f"{p.dtype=} != {t.dtype=}"
    2168        assert p.device == t.device, f"{p.device=} != {t.device=}"
    2169    if in_tree != in_tree2:
    2170        raise TypeError
    2171    f, out_tree_store = flatten_fn(f, in_tree, has_aux=has_aux)
    2172    jvp_ret = jvp_flat(
    2173        f,
    2174        primals_flat,
    2175        tangents_flat,
    2176        has_aux=has_aux,
    2177        global_data=global_data,
    2178        **static_args,
    2179    )
    2180    if has_aux:
    2181        (primals_out_flat, tangents_out_flat), aux = jvp_ret
    2182    else:
    2183        (primals_out_flat, tangents_out_flat) = jvp_ret
    2184    primals_out = tree_unflatten(out_tree_store(), primals_out_flat)
    2185    tangents_out = tree_unflatten(out_tree_store(), tangents_out_flat)
    2186    return ((primals_out, tangents_out), aux) if has_aux else (primals_out, tangents_out)
def jacfwd(f, x): View Source

    2189def jacfwd(f, x):
    2190    pushfwd = lambda v: jvp(f, (x,), (v,))[1]
    2191    vecs_in = backend.eye(math.prod(x.shape)).reshape(x.shape * 2)
    2192    return vmap(pushfwd, (0,))(vecs_in)
@contextmanager
def stash_trace(main: MainTrace): View Source

    2195@contextmanager
    2196def stash_trace(main: MainTrace):
    2197    global stashed_trace
    2198    prev_stashed_trace, stashed_trace = stashed_trace, main
    2199    try:
    2200        yield
    2201    finally:
    2202        stashed_trace = prev_stashed_trace
@contextmanager
def symbolic_run(): View Source

    2205@contextmanager
    2206def symbolic_run():
    2207    global trace_stack
    2208    level = len(trace_stack)
    2209    main = MainTrace(level, SymbolicRunTrace, global_data=None)
    2210    trace_stack += [main]
    2211    global stashed_trace
    2212    prev_stashed_trace, stashed_trace = stashed_trace, main
    2213    try:
    2214        yield
    2215    finally:
    2216        stashed_trace = prev_stashed_trace
    2217        trace_stack.pop()
def make_program(*args, **kwargs) -> Any: View Source

    172        def decorated_function(*args, **kwargs) -> Any:
    173            result = wrapper(*args, **kwargs)
    174            cache_info = wrapper.cache_info()
    175
    176            dblog(
    177                f"{fn.__name__}.{cache_info} {args.__hash__()}",
    178                enable=backend.LOG_LRU,
    179            )
    180            tb = "".join(traceback.format_list(traceback.extract_stack())[tb_start:tb_end]).replace("\n    ", ":\t") + "-" * 20 + "\n"
    181            dblog(f"{tb}", enable=backend.LOG_LRU)
    182
    183            return result
def vmap_program(*args, **kwargs) -> Any: View Source

    172        def decorated_function(*args, **kwargs) -> Any:
    173            result = wrapper(*args, **kwargs)
    174            cache_info = wrapper.cache_info()
    175
    176            dblog(
    177                f"{fn.__name__}.{cache_info} {args.__hash__()}",
    178                enable=backend.LOG_LRU,
    179            )
    180            tb = "".join(traceback.format_list(traceback.extract_stack())[tb_start:tb_end]).replace("\n    ", ":\t") + "-" * 20 + "\n"
    181            dblog(f"{tb}", enable=backend.LOG_LRU)
    182
    183            return result
def jvp_program(*args, **kwargs) -> Any: View Source

    172        def decorated_function(*args, **kwargs) -> Any:
    173            result = wrapper(*args, **kwargs)
    174            cache_info = wrapper.cache_info()
    175
    176            dblog(
    177                f"{fn.__name__}.{cache_info} {args.__hash__()}",
    178                enable=backend.LOG_LRU,
    179            )
    180            tb = "".join(traceback.format_list(traceback.extract_stack())[tb_start:tb_end]).replace("\n    ", ":\t") + "-" * 20 + "\n"
    181            dblog(f"{tb}", enable=backend.LOG_LRU)
    182
    183            return result
def partial_run_flat( f: Callable, pvals_in: List[PartialValue], has_aux, global_data=None) -> Tuple[Program, List[PartialValue], List[Any]]: View Source

    2275def partial_run_flat(
    2276    f: Callable, pvals_in: List["PartialValue"], has_aux, global_data=None
    2277) -> Tuple[Program, List["PartialValue"], List[Any]]:
    2278    with new_main_trace(PartialRunTrace, global_data) as main:
    2279        trace = PartialRunTrace(main)
    2280        tracers_in = [trace.new_arg(pval) for pval in pvals_in]
    2281        outs = f(*tracers_in)
    2282        if has_aux:
    2283            outs, aux = outs
    2284        tracers_out = [full_raise(trace, out) for out in outs]
    2285        pvals_out = [t.pval for t in tracers_out]
    2286        unk_tracers_in = [t for t in tracers_in if t.pval.is_unknown]
    2287        unk_tracers_out = [t for t in tracers_out if t.pval.is_unknown]
    2288        program, consts = tracers_to_program(unk_tracers_in, unk_tracers_out)
    2289
    2290    return (program, pvals_out, consts, aux) if has_aux else (program, pvals_out, consts)
def partial_run_program( program: Program, in_unknowns: List[bool], instantiate: Optional[List[bool]] = None) -> Tuple[Program, Program, List[bool], int]: View Source

    2293def partial_run_program(
    2294    program: Program,
    2295    in_unknowns: List[bool],
    2296    instantiate: Optional[List[bool]] = None,
    2297) -> Tuple[Program, Program, List[bool], int]:
    2298    env: Dict[Var, bool] = {}
    2299    residuals: Set[Var] = set()
    2300
    2301    def read(x: Atom) -> bool:
    2302        return type(x) is Var and env[x]
    2303
    2304    def write(unk: bool, v: Var) -> None:
    2305        env[v] = unk
    2306
    2307    instructions1, instructions2 = [], []
    2308    list_map(write, in_unknowns, program.in_binders)
    2309
    2310    for instruction in program.instructions:
    2311        unks_in = list_map(read, instruction.inputs)
    2312        (
    2313            instruction1,
    2314            instruction2,
    2315            unks_out,
    2316            res,
    2317        ) = instruction.op.partial_run_instruction(unks_in, instruction)
    2318        if instruction1 is not None:
    2319            instructions1 += [instruction1]
    2320        if instruction2 is not None:
    2321            instructions2 += [instruction2]
    2322        if res is not None:
    2323            residuals.update(res)
    2324        list_map(write, unks_out, instruction.out_binders)
    2325
    2326    out_unknowns = list_map(read, program.outs)
    2327    if instantiate is not None:
    2328        for v, uk, inst in zip(program.outs, out_unknowns, instantiate):
    2329            if inst and not uk:
    2330                if type(v) is Var:
    2331                    residuals.add(v)
    2332        out_unknowns = list_map(operator_py.or_, out_unknowns, instantiate)
    2333
    2334    residuals, num_res = list(residuals), len(residuals)
    2335    assert all(type(v) is Var for v in residuals), residuals
    2336
    2337    ins1, ins2 = partition_list(in_unknowns, program.in_binders)
    2338    outs1, outs2 = partition_list(out_unknowns, program.outs)
    2339
    2340    program1 = Program(
    2341        ins1,
    2342        instructions1,
    2343        outs1 + residuals,
    2344        0,
    2345        program.static_args,
    2346        f"{program.name}_partial1",
    2347    )
    2348    program2 = Program(
    2349        residuals + ins2,
    2350        instructions2,
    2351        outs2,
    2352        0,
    2353        program.static_args,
    2354        f"{program.name}_partial2",
    2355    )
    2356    typecheck_partial_run_program(program, in_unknowns, out_unknowns, program1, program2)
    2357
    2358    return program1, program2, out_unknowns, num_res
def typecheck_partial_run_program(program, in_unknowns, out_unknowns, program1, program2): View Source

    2361def typecheck_partial_run_program(program, in_unknowns, out_unknowns, program1, program2):
    2362    programty = typecheck_program(program)  # (a1,  a2) -> (b1, b2 )
    2363    program1ty = typecheck_program(program1)  #  a1       -> (b1, res)
    2364    program2ty = typecheck_program(program2)  # (res, a2) -> b2
    2365
    2366    a1, a2 = partition_list(in_unknowns, programty.in_types)
    2367    b1, b2 = partition_list(out_unknowns, programty.out_types)
    2368    b1_, res = split_list(program1ty.out_types, len(b1))
    2369    res_, a2_ = split_list(program2ty.in_types, len(res))
    2370    b2_ = program2ty.out_types
    2371
    2372    a1 = tuple(a1)
    2373    a2, a2_ = tuple(a2), tuple(a2_)
    2374    b1, b1_ = tuple(b1), tuple(b1_)
    2375    b2, b2_ = tuple(b2), tuple(b2_)
    2376    res, res_ = tuple(res), tuple(res_)
    2377
    2378    if program1ty.in_types != a1:
    2379        raise TypeError
    2380    if program2ty.out_types != b2:
    2381        raise TypeError
    2382    if b1 != b1_:
    2383        raise TypeError
    2384    if res != res_:
    2385        raise TypeError
    2386    if a2 != a2_:
    2387        raise TypeError
    2388    if b2 != b2_:
    2389        raise TypeError
def linearize_flat(f, *primals_in, has_aux): View Source

    2392def linearize_flat(f, *primals_in, has_aux):
    2393    pvals_in = [make_known_pval(x) for x in primals_in] + [make_unknown_pval(SymbolicTensor.like(get_symval(x))) for x in primals_in]
    2394
    2395    def f_jvp(*primals_tangents_in):
    2396        jvp_ret = jvp(f, *split_half(primals_tangents_in), has_aux=has_aux)
    2397        if has_aux:
    2398            (primals_out, tangents_out), aux = jvp_ret
    2399            return ((*primals_out, *tangents_out), aux)
    2400        else:
    2401            primals_out, tangents_out = jvp_ret
    2402            return (*primals_out, *tangents_out)
    2403
    2404    partial_run_flat_ret = partial_run_flat(f_jvp, pvals_in, has_aux)
    2405    if has_aux:
    2406        program, pvals_out, consts, aux = partial_run_flat_ret
    2407    else:
    2408        program, pvals_out, consts = partial_run_flat_ret
    2409    primal_pvals, _ = split_half(pvals_out)
    2410    assert all(pval.is_known for pval in primal_pvals)
    2411    primals_out = [pval.const for pval in primal_pvals]
    2412    f_lin = lambda *tangents: run_program(program, [*consts, *tangents])
    2413    return (primals_out, f_lin, aux) if has_aux else (primals_out, f_lin)
def linearize(f, *primals_in, has_aux=False): View Source

    2416def linearize(f, *primals_in, has_aux=False):
    2417    primals_in_flat, in_tree = tree_flatten(primals_in)
    2418    f, out_tree_store = flatten_fn(f, in_tree, has_aux=has_aux)
    2419    linearize_flat_ret = linearize_flat(f, *primals_in_flat, has_aux=has_aux)
    2420    if has_aux:
    2421        primals_out_flat, f_lin_flat, aux = linearize_flat_ret
    2422    else:
    2423        primals_out_flat, f_lin_flat = linearize_flat_ret
    2424
    2425    primals_out = tree_unflatten(out_tree_store(), primals_out_flat)
    2426
    2427    def f_lin(*tangents_in):
    2428        tangents_in_flat, in_tree2 = tree_flatten(tangents_in)
    2429        if in_tree != in_tree2:
    2430            raise TypeError
    2431        tangents_out_flat = f_lin_flat(*tangents_in_flat)
    2432        return tree_unflatten(out_tree_store(), tangents_out_flat)
    2433
    2434    return (primals_out, f_lin, aux) if has_aux else (primals_out, f_lin)
def tracers_to_program( tracers_in: List[PartialRunTraceTensor], tracers_out: List[PartialRunTraceTensor]): View Source

    2437def tracers_to_program(
    2438    tracers_in: List["PartialRunTraceTensor"],
    2439    tracers_out: List["PartialRunTraceTensor"],
    2440):
    2441    def tracer_parents(t: PartialRunTraceTensor) -> List[PartialRunTraceTensor]:
    2442        return t.draft.tracers_in if isinstance(t.draft, InstructionDraft) else []
    2443
    2444    def draft_to_instruction(tracer_to_var: Dict[int, Var], draft: InstructionDraft) -> Instruction:
    2445        inputs = [tracer_to_var[id(t)] for t in draft.tracers_in]
    2446        out_binders = [Var(symval) for symval in draft.symvals_out]
    2447        for t_ref, var in list_zip(draft.tracer_refs_out, out_binders):
    2448            if t_ref() is not None:
    2449                tracer_to_var[id(t_ref())] = var
    2450        return Instruction(draft.prim, inputs, draft.params, out_binders)
    2451
    2452    tracer_to_var: Dict[int, Var] = {id(t): Var(SymbolicTensor.like(t.symval)) for t in tracers_in}
    2453    constvar_to_val: Dict[int, Any] = {}
    2454    constid_to_var: Dict[int, Var] = {}
    2455    processed_instructions: Set[int] = set()
    2456    instructions: List[Instruction] = []
    2457    for t in toposort(tracers_out, tracer_parents):
    2458        if isinstance(t.draft, LambdaBindingDraft):
    2459            assert id(t) in set(list_map(id, tracers_in))
    2460        elif isinstance(t.draft, ConstDraft):
    2461            val = t.draft.val
    2462            var = constid_to_var.get(id(val))
    2463            if var is None:
    2464                symval = SymbolicTensor.like(get_symval(val))
    2465                var = constid_to_var[id(val)] = Var(symval)
    2466                constvar_to_val[var] = val
    2467            tracer_to_var[id(t)] = var
    2468        elif isinstance(t.draft, InstructionDraft):
    2469            if id(t.draft) not in processed_instructions:
    2470                instructions += [draft_to_instruction(tracer_to_var, t.draft)]
    2471                processed_instructions.add(id(t.draft))
    2472        else:
    2473            raise TypeError(t.draft)
    2474
    2475    constvars, constvals = unzip2(constvar_to_val.items())
    2476    in_binders = constvars + [tracer_to_var[id(t)] for t in tracers_in]
    2477    # out_vars = [tracer_to_var[id(t)] for t in tracers_out if id(t) in tracer_to_var]
    2478    out_vars = [tracer_to_var[id(t)] for t in tracers_out]
    2479    program = Program(tuple(in_binders), tuple(instructions), tuple(out_vars))
    2480    typecheck_program(program)
    2481    return program, constvals
def toposort(out_nodes: List[Any], parents: Callable[[Any], List[Any]]): View Source

    2484def toposort(out_nodes: List[Any], parents: Callable[[Any], List[Any]]):
    2485    def check_toposort(nodes: List[Any], parents: Callable[[Any], List[Any]]):
    2486        seen = set()
    2487        for node in nodes:
    2488            assert all(id(parent) in seen for parent in parents(node))
    2489            seen.add(id(node))
    2490
    2491    def remove_duplicates(lst):
    2492        seen = set()
    2493        return [x for x in lst if id(x) not in seen and not seen.add(id(x))]
    2494
    2495    if not out_nodes:
    2496        return []
    2497    out_nodes = remove_duplicates(out_nodes)
    2498
    2499    child_counts = {}
    2500    stack = list(out_nodes)
    2501    while stack:
    2502        node = stack.pop()
    2503        if id(node) in child_counts:
    2504            child_counts[id(node)] += 1
    2505        else:
    2506            child_counts[id(node)] = 1
    2507            stack.extend(parents(node))
    2508    for node in out_nodes:
    2509        child_counts[id(node)] -= 1
    2510
    2511    sorted_nodes = []
    2512    childless_nodes = [node for node in out_nodes if not child_counts[id(node)]]
    2513    while childless_nodes:
    2514        node = childless_nodes.pop()
    2515        sorted_nodes += [node]
    2516        for parent in parents(node):
    2517            if child_counts[id(parent)] == 1:
    2518                childless_nodes += [parent]
    2519            else:
    2520                child_counts[id(parent)] -= 1
    2521
    2522    sorted_nodes = sorted_nodes[::-1]
    2523    check_toposort(sorted_nodes, parents)
    2524    return sorted_nodes
def vjp_flat(f, *primals_in, has_aux=False, **static_args): View Source

    2527def vjp_flat(f, *primals_in, has_aux=False, **static_args):
    2528    pvals_in = [make_known_pval(x) for x in primals_in] + [make_unknown_pval(SymbolicTensor.like(get_symval(x))) for x in primals_in]
    2529    primal_pvals_in, tangent_pvals_in = split_half(pvals_in)
    2530    del primal_pvals_in
    2531
    2532    def f_jvp(*primals_tangents_in):
    2533        jvp_ret = jvp(
    2534            f,
    2535            *split_half(primals_tangents_in),
    2536            has_aux=has_aux,
    2537            global_data="vjp",
    2538            **static_args,
    2539        )
    2540        if has_aux:
    2541            ((primals_out, tangents_out), aux) = jvp_ret
    2542        else:
    2543            (primals_out, tangents_out) = jvp_ret
    2544        return ([*primals_out, *tangents_out], aux) if has_aux else [*primals_out, *tangents_out]
    2545
    2546    partial_run_flat_ret = partial_run_flat(f_jvp, pvals_in, has_aux, "vjp")
    2547    if has_aux:
    2548        program, pvals_out, consts, aux = partial_run_flat_ret
    2549    else:
    2550        program, pvals_out, consts = partial_run_flat_ret
    2551
    2552    primal_pvals, tangent_pvals = split_half(pvals_out)
    2553    del tangent_pvals
    2554    assert all(pval.is_known for pval in primal_pvals)
    2555    primals_out_flat = [pval.const for pval in primal_pvals]
    2556    transpose_inputs = consts + [UndefinedPrimal(t.symval) for t in tangent_pvals_in]
    2557
    2558    def f_vjp_flat(*cotangents):
    2559        # return backward_pass(program, transpose_inputs, cotangents)
    2560        undef_primals = tuple(isinstance(x, UndefinedPrimal) for x in transpose_inputs)
    2561        transposed_program, new_consts = transpose_program(program, undef_primals)
    2562        residuals, _ = partition_list(undef_primals, transpose_inputs)
    2563        outs = run_program(transposed_program, (*new_consts, *residuals, *cotangents))
    2564        return outs
    2565
    2566    return (primals_out_flat, f_vjp_flat, aux) if has_aux else (primals_out_flat, f_vjp_flat)
def vjp(f, *primals_in, has_aux=False, **static_args): View Source

    2569def vjp(f, *primals_in, has_aux=False, **static_args):
    2570    primals_in_flat, in_tree = tree_flatten(primals_in)
    2571    f, out_tree_store = flatten_fn(f, in_tree, has_aux=has_aux)
    2572    vjp_ret = vjp_flat(f, *primals_in_flat, has_aux=has_aux, **static_args)
    2573    if has_aux:
    2574        primals_out_flat, f_vjp_flat, aux = vjp_ret
    2575    else:
    2576        primals_out_flat, f_vjp_flat = vjp_ret
    2577    primals_out = tree_unflatten(out_tree_store(), primals_out_flat)
    2578
    2579    def f_vjp(*cotangents_out):
    2580        cotangents_out_flat, _ = tree_flatten(cotangents_out)
    2581        cotangents_in_flat = f_vjp_flat(*cotangents_out_flat)
    2582
    2583        return tree_unflatten(in_tree, cotangents_in_flat)
    2584
    2585    return (primals_out, f_vjp, aux) if has_aux else (primals_out, f_vjp)
NullCotangent = None
def backward_pass( program: Program, args: List[Any], cotangents: List[Any]) -> List[Any]: View Source

    2591def backward_pass(program: Program, args: List[Any], cotangents: List[Any]) -> List[Any]:
    2592    primal_env: Dict[Var, Any] = {}
    2593    ct_env: Dict[Var, Any] = {}
    2594
    2595    def read_primal(x: Atom) -> Any:
    2596        return primal_env.get(x, UndefinedPrimal(x.symval)) if type(x) is Var else x.val
    2597
    2598    def write_primal(v: Var, val: Any) -> None:
    2599        if type(val) is not UndefinedPrimal:
    2600            primal_env[v] = val
    2601
    2602    def read_cotangent(v: Var) -> Any:
    2603        return ct_env.pop(v, backend.zeros(v.symval.shape, v.symval.dtype))
    2604
    2605    def write_cotangent(x: Atom, ct: Any):
    2606        if type(x) is Var and ct is not NullCotangent:
    2607            ct_env[x] = (ct_env[x] + ct) if x in ct_env else ct
    2608
    2609    list_map(write_primal, program.in_binders, args)
    2610    list_map(write_cotangent, program.outs, cotangents)
    2611    for instruction in program.instructions[::-1]:
    2612        primals_in = list_map(read_primal, instruction.inputs)
    2613        cotangents_in = list_map(read_cotangent, instruction.out_binders)
    2614        inp, params = primals_in, instruction.params
    2615        cotangents_out = instruction.op.T(cotangents_in, *inp, **params)
    2616        list_map(write_cotangent, instruction.inputs, cotangents_out)
    2617
    2618    ret = [read_cotangent(v) for v, x in list_zip(program.in_binders, args) if isinstance(x, UndefinedPrimal)]
    2619    return ret
def transpose_program(*args, **kwargs) -> Any: View Source

    172        def decorated_function(*args, **kwargs) -> Any:
    173            result = wrapper(*args, **kwargs)
    174            cache_info = wrapper.cache_info()
    175
    176            dblog(
    177                f"{fn.__name__}.{cache_info} {args.__hash__()}",
    178                enable=backend.LOG_LRU,
    179            )
    180            tb = "".join(traceback.format_list(traceback.extract_stack())[tb_start:tb_end]).replace("\n    ", ":\t") + "-" * 20 + "\n"
    181            dblog(f"{tb}", enable=backend.LOG_LRU)
    182
    183            return result
def grad(f, argnums=(0,), argnames='', has_aux=False, return_value=False): View Source

    2640def grad(f, argnums=(0,), argnames="", has_aux=False, return_value=False):
    2641    f, rejit = (f, False) if not isinstance(f, jit) else (f.f, True)
    2642    if isinstance(argnums, int):
    2643        argnums = (argnums,)
    2644
    2645    def gfn(x, *xs, **static_args):
    2646        vjp_ret = vjp(f, x, *xs, has_aux=has_aux, **static_args)
    2647        if has_aux:
    2648            y, f_vjp, aux = vjp_ret
    2649        else:
    2650            y, f_vjp = vjp_ret
    2651        if np.shape(y) != ():
    2652            raise TypeError("grad output must be 0-dim scalar with shape ()")
    2653        gL_xs = f_vjp(backend.ones(()))
    2654        gL_xs = tuple(gL_xs[i] for i in argnums) if len(argnums) > 1 else gL_xs[argnums[0]]
    2655        if return_value:
    2656            return ((y, aux), gL_xs) if has_aux else (y, gL_xs)
    2657        else:
    2658            return (gL_xs, aux) if has_aux else gL_xs
    2659
    2660    return jit(gfn) if rejit else gfn
def value_and_grad(f, argnums=(0,), argnames='', has_aux=False): View Source

    2663def value_and_grad(f, argnums=(0,), argnames="", has_aux=False):
    2664    return grad(
    2665        f,
    2666        argnums=argnums,
    2667        argnames=argnames,
    2668        has_aux=has_aux,
    2669        return_value=True,
    2670    )
def jit_partial_run(trace, tracers, *, program): View Source

    2673def jit_partial_run(trace, tracers, *, program):
    2674    in_unknowns = [not t.pval.is_known for t in tracers]
    2675    program1, program2, out_unknowns, num_res = partial_run_program(program, in_unknowns)
    2676    known_tracers, unknown_tracers = partition_list(in_unknowns, tracers)
    2677    known_vals = [t.pval.const for t in known_tracers]
    2678    outs1_res = backend.jit_op(*known_vals, program=program)
    2679    outs1, res = split_list(outs1_res, len(program1.outs) - num_res)
    2680    res_tracers = [trace.instantiate_const(full_raise(trace, x)) for x in res]
    2681    outs2 = [PartialRunTraceTensor(trace, PartialValue.unknown(v.symval), None) for v in program2.outs]
    2682    draft = InstructionDraft(
    2683        backend.jit_op,
    2684        res_tracers + unknown_tracers,
    2685        dict(program=program2),
    2686        [v.symval for v in program2.outs],
    2687        map(weakref.ref, outs2),
    2688    )
    2689    for t in outs2:
    2690        t.draft = draft
    2691    return merge_lists(out_unknowns, outs1, outs2)
class jit: View Source

    2694class jit:
    2695    def __init__(self, f, static_argnames=(), name=None):
    2696        if isinstance(static_argnames, str):
    2697            static_argnames = tuple(static_argnames.split(" "))
    2698        assert type(static_argnames) is tuple and all(type(s) is str for s in static_argnames)
    2699        self.f = f
    2700        self.name = name if name is not None else self.f.__name__
    2701
    2702        self.static_argnames = static_argnames
    2703
    2704    @classmethod
    2705    def with_options(cls, static_argnames=(), name=None):
    2706        return partial(cls, static_argnames=static_argnames, name=name)
    2707
    2708    @classmethod
    2709    def get_jit_name(cls, args, static_args, prefix="jit", short=False):
    2710        name = f"{prefix}_"
    2711        if short:
    2712            static_args_tup = tuple(static_args.items())
    2713            ids = repr(hash((prefix, args, static_args_tup)))[-4:]
    2714            name = f"{prefix}_{ids}"
    2715        else:
    2716            for a in args:
    2717                name += f"shape_{a.shape}_dtype_{a.dtype.name}_"
    2718            for k, v in static_args.items():
    2719                name += f"{k}_{v}_"
    2720            name = name.replace("(", "_lp_")
    2721            name = name.replace(")", "_rp_")
    2722            name = name.replace(",", "_cm_")
    2723            name = name.replace(" ", "")
    2724            name = name.replace(".", "_dt_")
    2725
    2726        return name
    2727
    2728    def get_program(self, *args, **static_args):
    2729        sig = inspect.signature(self.f)
    2730        if all("*" not in repr(v) for v in sig.parameters.values()):
    2731            args_strs = [k for k, v in sig.parameters.items() if k != "self" and k not in self.static_argnames]
    2732            static_args_strs = [k for k, v in sig.parameters.items() if k != "self" and k in self.static_argnames]
    2733
    2734            if args:
    2735                if len(args) > len(args_strs):
    2736                    assert static_args_strs
    2737                    args, rest = args[: len(args_strs)], args[len(args_strs) :]
    2738                    new_static_args = {k: rest_arg for k, rest_arg in zip(static_args_strs, rest) if k not in static_args}
    2739                    static_args = {**new_static_args, **static_args}
    2740            else:
    2741                args = tuple([static_args[k] if k in static_args else arg for k, arg in zip(args_strs, args)])
    2742
    2743        symvals_in = tree_map(lambda x: SymbolicTensor.like(get_symval(x)), args)
    2744        static_args = tuple(static_args.items())
    2745        if self.name is None:
    2746            self.name = f"jit_{str(hash((self.f, symvals_in, static_args)))[-5:]}"
    2747        program, consts, out_tree = make_program(self.f, *symvals_in, static_args=static_args, name=self.name)
    2748        return program, consts, out_tree
    2749        program, consts, out_tree = make_program(
    2750            self.f,
    2751            *symvals_in,
    2752            static_args=tuple(static_args.items()),
    2753            name=self.name,
    2754        )
    2755        return program, consts, out_tree
    2756
    2757    def __call__(self, *args, **static_args):
    2758        program, consts, out_tree = self.get_program(*args, **static_args)
    2759        args, in_tree = tree_flatten(args)
    2760        # with Timing(f"JIT {self}"): outs = bind(backend.jit_op, *consts, *args, program=program)
    2761        outs = bind(backend.jit_op, *consts, *args, program=program)
    2762        return tree_unflatten(out_tree, outs)
    2763
    2764    def lower(self, *args, **static_args):
    2765        program, consts, out_tree = self.get_program(*args, **static_args)
    2766        args, in_tree = tree_flatten(args)
    2767        hashed_program = Hashed(program)
    2768        num_consts = program.num_consts
    2769        consts, args = args[:num_consts], args[num_consts:]
    2770        hashed_consts = tuple(map(Hashed, consts))
    2771        jit_output = backend.jit_program(hashed_program, hashed_consts)
    2772        return jit_output
    2773
    2774    def export(self, args, output_path, export_params=True, input_names=None, output_names=None, **kwargs):
    2775        if isinstance(args, Tensor):
    2776            args, static_args = (args,), dict()
    2777        elif not isinstance(args[-1], dict):
    2778            assert all(isinstance(a, Tensor) for a in args)
    2779            static_args = dict()
    2780        else:
    2781            args, static_args = args[:-1], args[-1]
    2782        assert isinstance(args, (tuple, list)) and isinstance(static_args, dict)
    2783        jit_output = self.lower(*args, **static_args)
    2784        backend.export(jit_output, output_path, export_params, input_names, output_names, **kwargs)
jit(f, static_argnames=(), name=None) View Source

    2695    def __init__(self, f, static_argnames=(), name=None):
    2696        if isinstance(static_argnames, str):
    2697            static_argnames = tuple(static_argnames.split(" "))
    2698        assert type(static_argnames) is tuple and all(type(s) is str for s in static_argnames)
    2699        self.f = f
    2700        self.name = name if name is not None else self.f.__name__
    2701
    2702        self.static_argnames = static_argnames
f
name
static_argnames
@classmethod
def with_options(cls, static_argnames=(), name=None): View Source

    2704    @classmethod
    2705    def with_options(cls, static_argnames=(), name=None):
    2706        return partial(cls, static_argnames=static_argnames, name=name)
@classmethod
def get_jit_name(cls, args, static_args, prefix='jit', short=False): View Source

    2708    @classmethod
    2709    def get_jit_name(cls, args, static_args, prefix="jit", short=False):
    2710        name = f"{prefix}_"
    2711        if short:
    2712            static_args_tup = tuple(static_args.items())
    2713            ids = repr(hash((prefix, args, static_args_tup)))[-4:]
    2714            name = f"{prefix}_{ids}"
    2715        else:
    2716            for a in args:
    2717                name += f"shape_{a.shape}_dtype_{a.dtype.name}_"
    2718            for k, v in static_args.items():
    2719                name += f"{k}_{v}_"
    2720            name = name.replace("(", "_lp_")
    2721            name = name.replace(")", "_rp_")
    2722            name = name.replace(",", "_cm_")
    2723            name = name.replace(" ", "")
    2724            name = name.replace(".", "_dt_")
    2725
    2726        return name
def get_program(self, *args, **static_args): View Source

    2728    def get_program(self, *args, **static_args):
    2729        sig = inspect.signature(self.f)
    2730        if all("*" not in repr(v) for v in sig.parameters.values()):
    2731            args_strs = [k for k, v in sig.parameters.items() if k != "self" and k not in self.static_argnames]
    2732            static_args_strs = [k for k, v in sig.parameters.items() if k != "self" and k in self.static_argnames]
    2733
    2734            if args:
    2735                if len(args) > len(args_strs):
    2736                    assert static_args_strs
    2737                    args, rest = args[: len(args_strs)], args[len(args_strs) :]
    2738                    new_static_args = {k: rest_arg for k, rest_arg in zip(static_args_strs, rest) if k not in static_args}
    2739                    static_args = {**new_static_args, **static_args}
    2740            else:
    2741                args = tuple([static_args[k] if k in static_args else arg for k, arg in zip(args_strs, args)])
    2742
    2743        symvals_in = tree_map(lambda x: SymbolicTensor.like(get_symval(x)), args)
    2744        static_args = tuple(static_args.items())
    2745        if self.name is None:
    2746            self.name = f"jit_{str(hash((self.f, symvals_in, static_args)))[-5:]}"
    2747        program, consts, out_tree = make_program(self.f, *symvals_in, static_args=static_args, name=self.name)
    2748        return program, consts, out_tree
    2749        program, consts, out_tree = make_program(
    2750            self.f,
    2751            *symvals_in,
    2752            static_args=tuple(static_args.items()),
    2753            name=self.name,
    2754        )
    2755        return program, consts, out_tree
def lower(self, *args, **static_args): View Source

    2764    def lower(self, *args, **static_args):
    2765        program, consts, out_tree = self.get_program(*args, **static_args)
    2766        args, in_tree = tree_flatten(args)
    2767        hashed_program = Hashed(program)
    2768        num_consts = program.num_consts
    2769        consts, args = args[:num_consts], args[num_consts:]
    2770        hashed_consts = tuple(map(Hashed, consts))
    2771        jit_output = backend.jit_program(hashed_program, hashed_consts)
    2772        return jit_output
def export( self, args, output_path, export_params=True, input_names=None, output_names=None, **kwargs): View Source

    2774    def export(self, args, output_path, export_params=True, input_names=None, output_names=None, **kwargs):
    2775        if isinstance(args, Tensor):
    2776            args, static_args = (args,), dict()
    2777        elif not isinstance(args[-1], dict):
    2778            assert all(isinstance(a, Tensor) for a in args)
    2779            static_args = dict()
    2780        else:
    2781            args, static_args = args[:-1], args[-1]
    2782        assert isinstance(args, (tuple, list)) and isinstance(static_args, dict)
    2783        jit_output = self.lower(*args, **static_args)
    2784        backend.export(jit_output, output_path, export_params, input_names, output_names, **kwargs)