search

fferflo / einx

Universal Tensor Operations in Einstein-Inspired Notation for Python.
https://einx.readthedocs.io/en/stable/
MIT License
328 stars 8 forks source link

MLX updates #12

Open altaic opened 3 months ago

altaic commented 3 months ago

FYI, einsum just landed in mlx along with python array api support. Also, thanks for einx!

https://github.com/ml-explore/mlx/pull/1269 https://github.com/ml-explore/mlx/pull/1289

fferflo commented 3 months ago

Hi, thanks for bringing this up! Support for mx.einsum (einx.dot) is already implemented, although not in a release version yet: https://github.com/fferflo/einx/commit/0ad7af3e675f6cd611cd0e38a7716f4eec81f66a It will be included in the next release of einx. (Using some other einx functions such as einx.get_at with mlx is currently limited to the use cases that are supported by mx.vmap)