Proposal: ❤️ Lovely Tensors + Lovely JAX

Lovely Tensors => Lovely JAX
Build Lovely JAX to be as good as Lovely Tensors

:heart: Lovely Tensors is a Python Library (that I wrote) that extends PyTorch with the ability to better summarize and visualize tensors in a human-friendly way. It is particularly useful during development/experimentation in Jupyter Notebooks but is in no way limited to Notebooks.

JAX is the new shiny low-level accelerated numerical computing library that is set to overtake PyTorch in the AI research field. I’m working on adding a Lovely Tensors counterpart in JAX.

Grant Deliverables:

  • Lovely JAX DeviceArray matches the functionality of Lovely Tensors
    • String repr
    • .rgb
    • .chans
    • .plot
  • Lovely Tensors extended to visualize/summarize NN layers (Linear, Conv, BatchNorm, etc.)

Squad

Squad Lead: xl0

  • Twitter @xl0xl0xl0
  • Alexey Zaytsev#6786
  • 0x45Fb23DFf8b6E73430f4e34688C820aA1134932F
4 Likes

You should get into touch with @_clashluke on twitter, i know he does alot of interesting JAX software.

2 Likes

Haha, yeah, he actually gave me the idea that a JAX version would be useful. :slight_smile:

This type of functionality would be super interesting to experiment with!

Yeah, I agree debugging ML models are pretty complicated task, and visualizing tensors is also tricky. This library can really help researchers and engineers in debugging ML codes.

Lovely tensors is awesome. I was just wondering how does .RGB visualization of a tensor look when its not of an image. Just because something like face is easy to catchup with but things might take a turn when its a RL model state.

My main intent is to see how to make debugging more humane in terms of how intuitive it is in a .plt plot of RGB values. Have you tried experimenting here…? I can try to contribute here

Hi. Thank you!

It all depends on the data. .rgb expects one of the dimensions to be of size 3 (RGB) or 4 (RGBA), but that’s the only hard limitation. And if it’s too restricting, you can use .chans, which works for any data that has at least 3 dimensions (Hmm, I can definitely move it to just 2).

If there are more dimensions before the H, W, C, the outputs are grouped automatically in a hierarchical way: lovely-tensors - Image grid

So, yes, please experiment, and let me know how it goes! I have not done any RL, so I don’t know what would be a useful way to visualize the common data there.