Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adding tensorboard tracking/visualization and nnx display example #118

Merged
merged 4 commits into from
Dec 3, 2024

Conversation

dillonroach
Copy link
Contributor

Starting from basic model/data in getting_started_with_jax_for_AI and modifying as appropriate to showcase tensorboard integration/tracking as well as basic nnx.display examples.

Copy link
Contributor

@vfdev-5 vfdev-5 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So, you do not want to showcase tbX in addition to pure tb/tf ?

docs/JAX_visualizing_models_metrics.md Outdated Show resolved Hide resolved
docs/JAX_visualizing_models_metrics.md Show resolved Hide resolved
@dillonroach
Copy link
Contributor Author

So, you do not want to showcase tbX in addition to pure tb/tf ?

Yea, I had mixed thoughts on this - I believe we were trying to stick as closely to the jax-ai-stack modules as possible and tbX wouldn't be in there, even if it is a good extra take on the API - I can get in a quick "and this is what it would look like in tbX" couple cells if we think it's an improvement and not taking away from the specific stack attention.

Copy link
Collaborator

@jakevdp jakevdp left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great! Just a couple comments below.

docs/JAX_visualizing_models_metrics.md Outdated Show resolved Hide resolved
docs/JAX_visualizing_models_metrics.md Outdated Show resolved Hide resolved
@dillonroach dillonroach requested a review from jakevdp December 3, 2024 06:06
@jakevdp
Copy link
Collaborator

jakevdp commented Dec 3, 2024

Looks good! We just need to resolve conflicts and then we can merge this!

@dillonroach
Copy link
Contributor Author

@jakevdp conflicts resolved

Copy link
Collaborator

@jakevdp jakevdp left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@jakevdp jakevdp merged commit 1998729 into jax-ml:main Dec 3, 2024
6 checks passed
windmaple pushed a commit to windmaple/jax-ai-stack that referenced this pull request Dec 5, 2024
…x-ml#118)

* adding tensorboard tracking/visualization and nnx display example

* Address Victor's initial feedback for information and links

* Addressing feedback - correcting links and final code cell
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants