mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 13:33:42 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
48 lines
2.0 KiB
Python
48 lines
2.0 KiB
Python
# NOTE: this is probably the wrong backend to work on, you want backend.py
|
|
# discussion about this in #pytorch-backend on our Discord
|
|
# this is using high level debug features from torch and we should integrate deeper
|
|
|
|
from tinygrad import Tensor
|
|
import torch, contextlib
|
|
from torch.utils._python_dispatch import TorchDispatchMode
|
|
from tinygrad.dtype import _from_torch_dtype
|
|
|
|
def empty_memory_format(size, dtype=None, layout=None, device=None, pin_memory=False, memory_format=None):
|
|
return TTensor(Tensor.empty(*size, dtype=_from_torch_dtype(dtype)))
|
|
|
|
# NOTE: if we have a way to change wrap/unwrap, these can be the same methods from backend.py
|
|
tiny_backend = {
|
|
"aten.empty.memory_format": empty_memory_format,
|
|
"aten.view.default": lambda x,sz: TTensor(x.tiny.reshape(sz)),
|
|
"aten.abs.default": lambda x: TTensor(x.tiny.abs()),
|
|
"aten.eq.Tensor": lambda x,y: TTensor(x.tiny == y.tiny),
|
|
"aten.bitwise_and.Tensor": lambda x,y: TTensor(x.tiny & y.tiny),
|
|
"aten.ne.Scalar": lambda x,y: TTensor(x.tiny != y),
|
|
"aten.mul.Tensor": lambda x,y: TTensor(x.tiny * y.tiny),
|
|
"aten.masked_select.default": lambda x,y: TTensor(Tensor(x.tiny.numpy()[y.tiny.numpy()])),
|
|
}
|
|
|
|
class TTensor(torch.Tensor):
|
|
tiny: Tensor
|
|
context = contextlib.nullcontext
|
|
|
|
@staticmethod
|
|
def __new__(cls, tiny, *args, **kwargs):
|
|
out = torch.Tensor._make_wrapper_subclass(cls, tiny.shape)
|
|
torch._C._set_throw_on_mutable_data_ptr(out)
|
|
out.tiny = tiny
|
|
return out
|
|
def __repr__(self): return super().__repr__(tensor_contents=f"{self.tiny}")
|
|
def __torch_dispatch__(cls, func, types, args, kwargs=None):
|
|
print(f"Dispatch Log: {func}(*{[type(x) for x in args]}, **{kwargs.keys()})")
|
|
#print(f"Dispatch Log: {func}(*{args}, **{kwargs})")
|
|
new_func = tiny_backend.get(str(func), None)
|
|
if new_func is None: raise NotImplementedError(f"add support for {func}")
|
|
return new_func(*args, **(kwargs or {}))
|
|
|
|
class Dispatcher(TorchDispatchMode): __torch_dispatch__ = TTensor.__torch_dispatch__
|
|
Dispatcher().__enter__()
|
|
|
|
if __name__ == "__main__":
|
|
a = torch.empty((4,), dtype=torch.int)
|