search

parrt / tensor-sensor

The goal of this library is to generate more helpful exception messages for matrix algebra expressions for numpy, pytorch, jax, tensorflow, keras, fastai.
https://github.com/parrt/tensor-sensor
MIT License
771 stars 39 forks source link

Supporting JAX #17

Closed ethanluoyc closed 3 years ago

ethanluoyc commented 3 years ago

Hi,

Thanks for the awesome library! This has really made my debugging life much easier.

Just a question. Is there any plan to support JAX? I think this can be similarly supported since the API of JAX almost looks identical to NumPy.

parrt commented 3 years ago

Heh, wow. jax looks really cool. I literally had not heard of this! Sure, I will poke around to see how this would work. parrt: check out https://sjmielke.com/jax-purify.htm

ethanluoyc commented 3 years ago

Thanks for the interest! It's getting more attention these days and I have found it to be quite useful in a broader sense of ML (probabilistic modeling). Please let me know if there is anything that I can help with.

parrt commented 3 years ago

I'm just about done with this feature request. Would you care to try it out? @ethanluoyc it is released so you can update if you like.

ethanluoyc commented 3 years ago

Wow this is absolutely amazing! I will give it a try and let you know how it works!

parrt commented 3 years ago

Lemme know what you think. it even works for weird stuff:

jnp.fft.fft(x, norm="something weird")
ethanluoyc commented 3 years ago

I just tried it out. It really works like magic. My work involves working with lots of broadcasting and reshaping so when something goes wrong it was usually very hard to understand what the problem is. Your work has really made things a lot simpler!

parrt commented 3 years ago

That's great to hear, @ethanluoyc ! It would be very helpful if you could advertise TensorSensor in the JAX community :) thank you!

ethanluoyc commented 3 years ago

Definitely!

On Dec 7, 2020, at 12:59 AM, Terence Parr notifications@github.com wrote:

That's great to hear, @ethanluoyc https://github.com/ethanluoyc ! It would be very helpful if you could advertise TensorSensor in the JAX community :) thank you!

— You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub https://github.com/parrt/tensor-sensor/issues/17#issuecomment-739604363, or unsubscribe https://github.com/notifications/unsubscribe-auth/ABOCZOBBBQCVAXIDQZN3AJ3STQSIXANCNFSM4UA4UMJA.