News for hackers. Stuff that matters.

A comment on Penzai: JAX research toolkit for building, editing, and visualizing neural nets. Parent comment

Re:Penzai: JAX research toolkit for building, editing, and v... by mccoyb on Sunday April 21, 2024 @02:10PM (#40107859)
That's cool -- but wouldn't it be more constructive to discuss "the ideas" in this package anyways?

For instance, it would be interesting to discern if the design of PyTorch (and their modules) preclude or admit the same sort of visualization tooling? If you have expertise in PyTorch, perhaps you could help answer this sort of question?

JAX's Pytrees are like "immutable structs, with array leaves" -- does PyTorch have a similar concept?

[ Parent | on HN ]
Re:Penzai: JAX research toolkit for building, editing, and v... by fpgamlirfanboy on Sunday April 21, 2024 @02:21PM (#40107946)
> does PyTorch have a similar concept

of course https://github.com/pytorch/pytorch/blob/main/torch/utils/_py...

[ Parent | on HN ]
Re:Penzai: JAX research toolkit for building, editing, and v... by ein0p on Sunday April 21, 2024 @02:13PM (#40107881)
Idk if you need that immutability actually. You could probably reconstruct enough to do this kind of viz from the autograd graph, or capture the graph and intermediates in the forward pass using hooks. My hunch is it should be doable.
[ Parent | on HN ]