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

Add automatic parallelization JAX example #4973

Merged
merged 16 commits into from
Aug 4, 2023

Conversation

awolant
Copy link
Contributor

@awolant awolant commented Aug 2, 2023

Depends on #4969

Category:

New feature

Description:

Adds example of automatic parallelization with JAX and DALI.

Additional information:

Affected modules and functionalities:

JAX plugin docs.

Key points relevant for the review:

Is it understandable? Spelling and syntax?

Tests:

  • Existing tests apply
  • New tests added
    • Python tests
    • GTests
    • Benchmark
    • Other
  • N/A

Checklist

Documentation

  • Existing documentation applies
  • Documentation updated
    • Docstring
    • Doxygen
    • RST
    • Jupyter
    • Other
  • N/A

DALI team only

Requirements

  • Implements new requirements
  • Affects existing requirements
  • N/A

REQ IDs: N/A

JIRA TASK: DALI-3558

Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
@review-notebook-app
Copy link

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
@awolant awolant marked this pull request as ready for review August 3, 2023 09:38
Copy link
Contributor Author

Choose a reason for hiding this comment

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

This file is part of #4969

Copy link
Contributor Author

Choose a reason for hiding this comment

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

This file is part of #4969

@awolant awolant changed the title [WIP] Add auto matic parallelization JAX example Add auto matic parallelization JAX example Aug 3, 2023
@awolant awolant changed the title Add auto matic parallelization JAX example Add automatic parallelization JAX example Aug 3, 2023
Signed-off-by: Albert Wolant <awolant@nvidia.com>
@@ -27,7 +27,7 @@
},
Copy link
Contributor

@mzient mzient Aug 3, 2023

Choose a reason for hiding this comment

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

(single GPU training example)[jax-basic_example.ipynb]

Is it some kind of markup for Sphinx or sth?


Reply via ReviewNB

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Fixed

Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
Signed-off-by: Albert Wolant <awolant@nvidia.com>
@awolant
Copy link
Contributor Author

awolant commented Aug 3, 2023

!build

@dali-automaton
Copy link
Collaborator

CI MESSAGE: [9229567]: BUILD STARTED

@dali-automaton
Copy link
Collaborator

CI MESSAGE: [9230402]: BUILD STARTED

@dali-automaton
Copy link
Collaborator

CI MESSAGE: [9229567]: BUILD PASSED

@dali-automaton
Copy link
Collaborator

CI MESSAGE: [9230402]: BUILD PASSED

@awolant awolant merged commit 9a82fcc into NVIDIA:main Aug 4, 2023
5 checks passed
JanuszL pushed a commit to JanuszL/DALI that referenced this pull request Oct 13, 2023
Adds example of automatic parallelization with JAX and DALI.

Signed-off-by: Albert Wolant <awolant@nvidia.com>
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.

None yet

4 participants