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 general system class #425

Merged
merged 32 commits into from
Mar 9, 2022

Conversation

arnupretorius
Copy link
Collaborator

@arnupretorius arnupretorius commented Feb 23, 2022

What?

General system class for Jax-based Mava systems.

Why?

The new system class supports a more general and flexible build system using callback components.

How?

Create a system class with functions for specifying a system design using components and adding, updating and configuring system component hyperparameters.

Extra

I'm experiencing a tricky mypy issue that I haven't been able to solve. It has to do with the Builder attributes not being picked up even though the tests seem to indicate that the system and builder interact as expected. This is also the case with the CallbackHookMixin. @KaleabTessera perhaps you might have some ideas once you have looked at the code?

closes #406

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.

This is really great!!! 🔥 Just made some minor comments :)

mava/systems/jax/config.py Show resolved Hide resolved
mava/systems/jax/system.py Show resolved Hide resolved
mava/systems/jax/system.py Show resolved Hide resolved
mava/systems/jax/system.py Outdated Show resolved Hide resolved
KaleabTessera
KaleabTessera previously approved these changes Mar 8, 2022
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 🦾

AsadJeewa
AsadJeewa previously approved these changes Mar 8, 2022
mava/systems/jax/system.py Show resolved Hide resolved
mava/systems/jax/config_test.py Show resolved Hide resolved
mava/systems/jax/config_test.py Show resolved Hide resolved
mava/systems/jax/system.py Outdated Show resolved Hide resolved
AsadJeewa
AsadJeewa previously approved these changes Mar 9, 2022
mava/systems/jax/config_test.py Show resolved Hide resolved
RuanJohn
RuanJohn previously approved these changes Mar 9, 2022
@arnupretorius arnupretorius dismissed stale reviews from RuanJohn and AsadJeewa via 62e8d28 March 9, 2022 10:38
@AsadJeewa AsadJeewa self-requested a review March 9, 2022 11:02
@arnupretorius arnupretorius merged commit 346d6c4 into develop Mar 9, 2022
@arnupretorius arnupretorius deleted the feature/jax-general-system-class branch March 9, 2022 11:16
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/XL
Projects
None yet
Development

Successfully merging this pull request may close these issues.

[FEATURE] General system class for Jax systems
4 participants