search

jonasrauber / eagerpy

PyTorch, TensorFlow, JAX and NumPy — all of them natively using the same code
https://eagerpy.jonasrauber.de
MIT License
693 stars 39 forks source link

Support for @ operator ? #35

Closed eserie closed 3 years ago

eserie commented 3 years ago

It could be great to have support for the @ operator.

Given two tensors:

a = ep.astensor(np.random.normal(size=(2, 10)))
b = ep.astensor(np.random.normal(size=(10, 3)))

The expression:

c = a.matmul(b)

could be written as:

c = a @ b
jonasrauber commented 3 years ago

Agree. I think all we have to do is to overwrite __matmul__ to delegate to matmul.

eserie commented 3 years ago

Ok great! Do you want I propose a PR for that?

jonasrauber commented 3 years ago

Feel free to do so.

eserie commented 3 years ago

I created the PR https://github.com/jonasrauber/eagerpy/pull/37#.

jonasrauber commented 3 years ago

merged #37, so closing this :)

eserie commented 3 years ago

Thanks!