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']