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

Update _combining.py changing examples to use optax.multi_transform, not optax.partition #1205

Open
wants to merge 3 commits into
base: main
Choose a base branch
from

Conversation

ywrt
Copy link

@ywrt ywrt commented Feb 23, 2025

Change mention of optax.partition to optax.multi_transform as optax.partition no longer exists.

Change mention of optax.partition to optax.multi_transform as optax.partition no longer exists.
Copy link

google-cla bot commented Feb 23, 2025

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@rdyro
Copy link
Collaborator

rdyro commented Feb 23, 2025

This looks great!

Could you also update:

  • optax/transforms/_combining_test.py to use multi_transform
  • optax/transforms/__init__.py to export multi_transform in addition to the old partition (essentially adding multi_transform = partition and adding multi_transform to __all__)

Thank you!

ywrt added 2 commits February 23, 2025 21:26
Rename parition -> multi_transform, PartitionState -> MultiTransformState,  rename the tests, and amend tests to use new name.
@ywrt
Copy link
Author

ywrt commented Feb 23, 2025

I fully renamed partition to multi_transform, and added partition/PartitionState aliases for backward compatibility.

@rdyro
Copy link
Collaborator

rdyro commented Feb 24, 2025

It looks great! Just a minor comment and let's merge this, thank you!

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.

2 participants