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

Feature/jax abstract builder class #433

Merged
merged 10 commits into from
Mar 2, 2022

Conversation

arnupretorius
Copy link
Collaborator

@arnupretorius arnupretorius commented Feb 24, 2022

What?

Abstract builder class for jax-based Mava systems.

Why?

To inform the builder for Jax systems through inheritance.

How?

Create class with the required functions, including building a data and parameter server, executor and trainer process.

Extra

@DriesSmit I decided to remove the evaluator as a separate build function since I think we can just use the executor and specify that it is an evaluator process using executor_id="evaluator". Let me know what you think 🙂

closes #429

mava/core_jax.py Outdated Show resolved Hide resolved
Copy link
Contributor

@DriesSmit DriesSmit left a comment

Choose a reason for hiding this comment

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

I really like this setup! Really cool work, thanks @arnupretorius 🔥 🔥 Just see my comment.

DriesSmit
DriesSmit previously approved these changes Feb 24, 2022
@pull-request-size pull-request-size bot added size/M and removed size/L labels Feb 25, 2022
Copy link
Contributor

@AsadJeewa AsadJeewa left a comment

Choose a reason for hiding this comment

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

Looks good. Difficult to review PRs with abstract classes but all seems in order and keen to see concrete implementations. Thanks

Copy link
Contributor

@KaleabTessera KaleabTessera left a comment

Choose a reason for hiding this comment

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

Thanks @arnupretorius , just some comments/clarifications!

mava/core_jax.py Show resolved Hide resolved
mava/core_jax.py Show resolved Hide resolved
mava/core_jax.py Show resolved Hide resolved
mava/core_jax.py Show resolved Hide resolved
@arnupretorius arnupretorius merged commit d846983 into develop Mar 2, 2022
@arnupretorius arnupretorius deleted the feature/jax-abstract-builder-class branch March 2, 2022 10:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request size/M
Projects
None yet
Development

Successfully merging this pull request may close these issues.

[FEATURE] Abstract builder class for Jax-based systems
5 participants