search

arogozhnikov / einops

Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)
https://einops.rocks
MIT License
8.41k stars 350 forks source link

Tinygrad support #297

Closed blueridanus closed 9 months ago

blueridanus commented 10 months ago

Might want to wait a bit on: https://github.com/tinygrad/tinygrad/pull/2686 Done.

blueridanus commented 10 months ago

https://github.com/tinygrad/tinygrad/pull/2686 merged, ready for review 🖖

kroggen commented 9 months ago

@arogozhnikov Is it possible to merge this?

arogozhnikov commented 9 months ago

I wish you folks implemented array api instead - and got einops/sklearn support for free.

E.g. in eindex I won't spend my time implementing backends for different frameworks - that's too much hassle

geohot commented 9 months ago

@arogozhnikov We are interested in being array API compatible. Totally understand not writing special backend code

Is there a good compliance test for this? How far off are we?

arogozhnikov commented 9 months ago

@geohot Array API comes with a test suite to verify compatibility https://data-apis.org/array-api/2022.12/verification_test_suite.html

to see how far you off, start with this additional method in Tensor:

def __array_namespace__(self, api_version: Optional[str] = None) -> Any:
    return Tensor # tensor class, as it seems to contain most of relevant methods.

I expect that due to divergence it would be easier to have another namespace like tinygrad.array_api with corresponding methods.

In a perfect world __array_namespace__ would return tinygrad itself, but I see tinygrad keeps minimal exports in root namespace.