Trouble understading error message : ValueError: Can't call compact methods on unbound modules

[LysSanzMoreta]

Hi!

I am following the Seq2Seq example https://github.com/google/flax/blob/master/examples/seq2seq/train.py to implement my own little Decoder and I am having some issues with how to place the classes and decorators appropiately, so any help is appreciated. I do not even need to do time series (although the example is a time series), in this case I would like to feed the entire batch [batch_size,max_len,input_size] to the GRU and then obtain the logits [batch_size,max_len,logits_size]. I was able to implement this in a previous version of flax, but I struggle with the new framework. I hope it's not a problem of compatibility between numpyro and flax.
I have

• jaxlib 1.59+cuda111
• flax: 0.3.0
• numpyro: 0.5.0

My little example is as follows:

class GRU(nn.Module):
  teacher_force: bool
  aa_prob : float
  @functools.partial(
      nn.transforms.scan,
      variable_broadcast='params', #I am struggling to understand the role of these, I only see that they are partial inputs to scan
      split_rngs={'params': False})
  @nn.compact
  def __call__(self, carry,x):
    """Carry=rng, Hidden state
    x = Input"""
    rng, gru_state= carry
    carry_rng, categorical_rng = jax.random.split(rng, 2)
    gru_state, y = nn.GRUCell()(gru_state, x)
    logits = nn.Dense(features=self.aa_prob)(y)
    return (carry_rng, gru_state), (logits)

class DecoderGRU(nn.Module):
  """GRU decoder seq2seq """
  teacher_force: bool
  aa_prob :float
  hidden_dim:float
  @nn.compact
  def __call__(self, inputs):
      gru = GRU(teacher_force=self.teacher_force,aa_prob = self.aa_prob)
      rng = jax.random.PRNGKey(1)
      init_state = nn.GRUCell.initialize_carry(jax.random.PRNGKey(0), (), self.hidden_dim)
      init_carry = (rng, init_state) #changed from self.make_rng("gru")
      input = jax.lax.slice_in_dim(inputs, 0, 1)[0] #remove?
      _, (logits) = gru(init_carry, input)
      return logits

And then I call it from another script as such. I have tried many things but none worked, I put the simplest manner

def model(*args):
    decoder = DecoderGRU(teacher_force=False,aa_prob =args.aa_prob,hidden_dim=args.gru_hidden_dim)
    logits = decoder(decoder_inputs) #simplest approach, decoder_inputs has size [batch_size, max_len,input_size)
    

My error message

Traceback (most recent call last):
File "/home/lys/.../Draupnir_Bayes.py", line 279, in
main(Dataset_train, Dataset_test, matrix)
File "/home/lys/.../Draupnir_Bayes.py", line 198, in main
mcmc.run(rng,rng_init,Dataset_train,sequences, angles,matrix,args)
File "/home/lys/anaconda3/lib/python3.7/site-packages/numpyro/infer/mcmc.py", line 499, in run
states_flat, last_state = partial_map_fn(map_args)
File "/home/lys/anaconda3/lib/python3.7/site-packages/numpyro/infer/mcmc.py", line 328, in _single_chain_mcmc
model_args=args, model_kwargs=kwargs)
File "/home/lys/anaconda3/lib/python3.7/site-packages/numpyro/infer/hmc.py", line 468, in init
init_params = self._init_state(rng_key_init_model, model_args, model_kwargs, init_params)
File "/home/lys/anaconda3/lib/python3.7/site-packages/numpyro/infer/hmc.py", line 430, in _init_state
forward_mode_differentiation=self._forward_mode_differentiation)
File "/home/lys/anaconda3/lib/python3.7/site-packages/numpyro/infer/util.py", line 431, in initialize_model
substituted_model, model_args, model_kwargs)
File "/home/lys/anaconda3/lib/python3.7/site-packages/numpyro/infer/util.py", line 278, in _get_model_transforms
model_trace = trace(model).get_trace(*model_args, **model_kwargs)
File "/home/lys/anaconda3/lib/python3.7/site-packages/numpyro/handlers.py", line 161, in get_trace
self(*args, **kwargs)
File "/home/lys/anaconda3/lib/python3.7/site-packages/numpyro/primitives.py", line 81, in call
return self.fn(*args, **kwargs)
File "/home/lys/anaconda3/lib/python3.7/site-packages/numpyro/primitives.py", line 81, in call
return self.fn(*args, **kwargs)
File "/home/lys/anaconda3/lib/python3.7/site-packages/numpyro/primitives.py", line 81, in call
return self.fn(*args, **kwargs)
File "/home/lys/Dropbox/PhD/DRAUPNIR/Draupnir_Bayes_models.py", line 41, in model1
logits = decoder(decoder_inputs)
File "/home/lys/anaconda3/lib/python3.7/site-packages/flax/linen/module.py", line 214, in wrapped_module_method
raise ValueError("Can't call compact methods on unbound modules")
ValueError: Can't call compact methods on unbound modules

I would very much appreciate any input on how to structure my code. Thanks! :)

[avital]

We absolutely need to improve this error message (being tracked at #893 -- I think the error may have changed, but it's the same issue). In short, you have to pass the variables (parameters and other state) via Module.apply (or use Module.init to get initial variables)

In your case, instead of

logits = decoder(decoder_inputs)

you should do:

variables = {"params": ..., ...any other state you have...}
logits = decoder.apply(variables, decoder_inputs)

Constructing a module instance does not give it variables, this is the same as pre-Linen Flax.

[LysSanzMoreta]

Ops! Wait, I need more input, sorry. So I am not sure how many states and how do I have to initialize them, so I tried:

key = random.PRNGKey(1)
variables_dict = {"params": key,"gru":key,"kernel":key} #I am afraid I am not sure how many things do I hav to initialize here
variables = decoder.init(variables_dict,decoder_input)["params"]
# apply the NN to the variables + input to get output
logits = decoder.apply(variables, decoder_input)

And then I get this error:

File "/home/lys/anaconda3/lib/python3.7/site-packages/flax/core/scope.py", line 546, in param
raise ValueError(f'No paramater named "{name}" exists in "{self.path_text}".')
ValueError: No paramater named "kernel" exists in "/scan(GRU_0)/GRUCell_0/ir".

I guess I need to initialize "kernel" somehow, but I tried to add it to the variables_dict and it did not work. I also tried to do initialize "dense", "inner" etc. I apologize for not understanding this new structure. Thanks again

[avital]

The first argument to init is not a variables dict, it's the random keys used. So unless you're doing something special,

variables = decoder.init({"params": key}, decoder_input)

should work (in fact decoder.init(key, decoder_input) also works). But don't access just the parameters, keep the full variable dict. So this should work:

key = random.PRNGKey(1)
variables = decoder.init(key}, decoder_input)
logits = decoder.apply(variables, decoder_input)

[LysSanzMoreta]

Oh, ok, thanks! seems to work now :)