Inconsistent shapes between value and initializer for parameter: FlaxGPT2LMHeadModel
Author: 0xnakul
Pull request: no

Hello!
I was trying to finetune GPT-2 medium model through flax on a custom (tokenized) dataset and I encoountered this error: `Inconsistent shapes between value and initializer for parameter "scale" in "/transformer/ln_f": (1024,), (0,)`

Edit: The whole traceback is quite long and is reported [here](https://pastebin.com/UMa3BbxP). A very short version is mentioned at the end here.

I'm using a PyTorch Dataset(with 1024 tokens per batch) and DataLoader(`batch_size=64`) with a `numpy_collate` function as mentioned at https://jax.readthedocs.io/en/latest/notebooks/Neural_Network_and_Data_Loading.html and then I'm yielding a "superbatch" of shape (8, 64, 1024) for multi-tpu using a custom function.

I'm using pre-trained gpt-2 tokenizer along with FlaxGPT2LMHeadModel.

Here is the code for the training loop:
```
for epoch in tqdm(range(1, num_epochs + 1), desc=f"Epoch ...", position=0, leave=True):
	    rng, input_rng = jax.random.split(rng)

    # -- Train --
    train_loader = make_superbatch()
    with tqdm(total=len(script_dataset), desc="Training...", leave=False) as progress_bar_train:
        for model_inputs in train_loader:
            # Model forward
            state, train_metric, dropout_rngs = parallel_train_step(state, model_inputs, dropout_rngs)

            progress_bar_train.update(1)

        progress_bar_train.write(
              f"Train... ({epoch}/{num_epochs} | Loss: {round(train_metric['loss'].mean(), 3)}, Learning Rate: {round(train_metric['learning_rate'].mean(), 6)})"
        )
```
Here is the error that I'm encountering and the whole traceback:
```
Epoch ...:   0%|          | 0/10 [00:00<?, ?it/s]
Training...:   0%|          | 0/1470930 [00:00<?, ?it/s]
                                                        

---------------------------------------------------------------------------

UnfilteredStackTrace                      Traceback (most recent call last)

<ipython-input-29-5c831c772fc6> in <module>()
      8             # Model forward
----> 9             state, train_metric, dropout_rngs = parallel_train_step(state, model_inputs, dropout_rngs)
     10 

47 frames

UnfilteredStackTrace: flax.errors.ScopeParamShapeError: Inconsistent shapes between value and initializer for parameter "scale" in "/transformer/ln_f": (1024,), (0,). (https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.ScopeParamShapeError)

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

The above exception was the direct cause of the following exception:

ScopeParamShapeError                      Traceback (most recent call last)

/usr/local/lib/python3.7/dist-packages/flax/core/scope.py in param(self, name, init_fn, *init_args)
    618         if jnp.shape(val) != jnp.shape(abs_val):
    619           raise errors.ScopeParamShapeError(name, self.path_text, 
--> 620               jnp.shape(val), jnp.shape(abs_val))
    621     else:
    622       if not self.is_mutable_collection('params'):

ScopeParamShapeError: Inconsistent shapes between value and initializer for parameter "scale" in "/transformer/ln_f": (1024,), (0,). (https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.ScopeParamShapeError)
```


I'm guessing this `1024` comes from the number of tokens per batch. How do I resolve this error?
Any help would be much appreciated.

Thank You.


Comments:
I don't exactly know how but this issue went away on its own after appearing unexpectedly out of nowhere. Closing the issue. 

> I don't exactly know how but this issue went away on its own after appearing unexpectedly out of nowhere. Closing the issue.

Dear thisis-nkul, have you sovled this problem? How?