Merge pull request #3343 from chiamp:haiku_upgrade

PiperOrigin-RevId: 565796217
google · Sep 15, 2023 · 5d846a5 · 5d846a5
2 parents 654ae1a + 089f237
commit 5d846a5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/docs/guides/haiku_migration_guide.rst b/docs/guides/haiku_migration_guide.rst
@@ -891,7 +891,7 @@ The Flax team recommends a more Module-centric approach that uses `__call__` to
 
   model = hk.transform_with_state(forward)
 
-  params, state = model.init(PRNGKey(0), jax.numpy.ones((1, 64)))
+  params, state = model.init(random.key(0), jax.numpy.ones((1, 64)))
 
   ---
 
@@ -906,5 +906,5 @@ The Flax team recommends a more Module-centric approach that uses `__call__` to
       return output
 
   model = FooModule()
-  variables = model.init(PRNGKey(0), jax.numpy.ones((1, 64)))
+  variables = model.init(random.key(0), jax.numpy.ones((1, 64)))
   params, counter = variables['params'], variables['counter']