FLIP: rm 'shape' from Module.param call signature about flax HOT 8 CLOSED

I agree the shape argument is arbitrary. It also implies that we can only return a single tensor which is also arbitrary.

Note that Flax is using the initializers from jax.nn.initializers so that create some difficulty when we want to reuse those.

An alternative that is less invasive is the following:

Change the current param api:

def param(self, shape, init_fn):
  ...
  init_value = init_fn(rng, shape)

to:

def param(self[, shape=None], init_fn):
  if shape is None:
    init_value = init_fn(rng):
  else:
    init_value = init_fn(rng, shape)

from flax.

Isn't

def param(self, shape=None, init_fn):

a SyntaxError? Just wondering if there's a way around that...

from flax.

Note that Flax is using the initializers from jax.nn.initializers so that create some difficulty when we want to reuse those.

I think it's pretty straightforward to wrap the JAX initializers though.

from flax.

Isn't

def param(self, shape=None, init_fn):

a SyntaxError? Just wondering if there's a way around that...

Yes you have to implement some manual argument unpacking to get that working. With the brackets I meant the param method is basically overloaded so you have either

param(name, shape, init_fn)
or
param(name, init_fn)

we could also decide to give the second option a different name.

from flax.

from flax.

If both Flax and Haiku want a different jax.nn.initializers API we should maybe switch it...

from flax.

@jekbradbury - oh, do haiku peeps also want a different initializer api?

from flax.

This has been resolved (with an improved shape check) in #530

from flax.