Argument init in get_parameter is not optional

[hylkedonker]

Hi,

According to the docs, the init argument of get_parameter is optional, while in reality it raises the error:
ValueError: Initializer must be specified.
(See line 496 in base.py.)

Example
An example where the init=None might occur is if the initialisation is done outside Haiku.
For example:

import haiku as hk
import jax.numpy as jnp

@hk.without_apply_rng
@hk.transform
def foo(x):
    w = hk.get_parameter("w", [1], init=None)
    return x + w

# Initialise params outside haiku, without `foo.init`.
params = {'~': {'w': jnp.array([2.], dtype=jnp.float32)}}

x = jnp.array([1])
foo.apply(params, x)

Kind regards,

Hylke

[tomhennigan]

Thank you Hylke, I'll fix this for the next release of Haiku. For now you can work around this by passing in a throwing initializer:

def throw_if_run(shape, dtype):
  del shape, dtype
  raise ValueError("Initializer must be specified.")

w = hk.get_parameter("w", [], init=throw_if_run)

[hylkedonker]

Great, thanks a lot!