search

jax-ml / jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
http://jax.readthedocs.io/
Apache License 2.0
30.38k stars 2.79k forks source link

Integration with TensorBoard #135

Closed alexbw closed 5 years ago

alexbw commented 5 years ago

This could be a simple example of how to do it, since PyTorch users have already extracted it from TensorFlow as a whole (https://github.com/lanpa/tensorboardX, but that's just one example)

hawkinsp commented 5 years ago

I'm wondering if we still want to do this in Jax itself.

Trax already has a TensorBoard integration (jaxboard): https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/trax/jaxboard.py

mattjj commented 5 years ago

I think we do not want to do this in JAX itself. We're starting to point more and more to trax as an example, and this is one more reason to do it. We could point to the TensorBoard integration specifically, but since trax mentions it prominently in its readme, I don't think that's necessary.