Skip to content
This repository has been archived by the owner on Feb 26, 2023. It is now read-only.

Bumps flax to 0.4.0 #60

Merged
merged 3 commits into from Feb 4, 2022
Merged

Bumps flax to 0.4.0 #60

merged 3 commits into from Feb 4, 2022

Conversation

ptigwe
Copy link
Contributor

@ptigwe ptigwe commented Feb 4, 2022

Updates flax to the most recent version. This currently breaks the
current implementation and way in which rng keys of dropout is being
handled.

Currently have disabled one of the dropout equivalence tests as I am not
fully aware if there is a method of directly affecting the value of
next_key within a treex module.

ptigwe and others added 3 commits February 4, 2022 22:56
Updates flax to the most recent version. This currently breaks the
current implementation and way in which rng keys of  dropout is being
handled.

Currently have disabled one of the dropout equivalence tests as I am not
fully aware if there is a method of directly affecting the value of
`next_key` within a treex module.
@cgarciae
Copy link
Owner

cgarciae commented Feb 4, 2022

Hey @ptigwe! Thanks a lot for this :)

Currently have disabled one of the dropout equivalence tests as I am not
fully aware if there is a method of directly affecting the value of
next_key within a treex module.

Yeah, trying to simulate the rng key splits is tricky (I've fought with it in the past). Managed to figure it out though so I pushed the fix + also applied pre-commit which was failing.

I think this is good to go once CI passes.

@cgarciae cgarciae merged commit 4b9bdac into cgarciae:master Feb 4, 2022
@ptigwe ptigwe deleted the flax_update branch February 5, 2022 01:13
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants