Using resnet-18 in flax
Author: aagha6
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.52.4
- Platform: macOS-15.5-arm64-arm-64bit
- Python version: 3.9.21
- Huggingface_hub version: 0.32.4
- Safetensors version: 0.5.3
- Accelerate version: not installed
- Accelerate config: not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.6.0 (False)
- Tensorflow version (GPU?): 2.18.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.8.5 (cpu)
- Jax version: 0.4.30
- JaxLib version: 0.4.30
- Using distributed or parallel set-up in script?: No

### Who can help?

@amyeroberts @Rocketknight1  

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

I'm trying to use resnet-18 in flax, for some reason, it seems like the normalization layers can't be loaded. 

I load the model with the command `model = FlaxResNetModel.from_pretrained("microsoft/resnet-18", dtype=jnp.float32)`.

I get this error, It's also performing much worse as a backbone for my task in comparsion with `resnet-26`. 

`Some weights of FlaxResNetModel were not initialized from the model checkpoint at microsoft/resnet-18 and are newly initialized: {('batch_stats', 'encoder', 'stages', '2', 'layers', '0', 'layer', 'layer_0', 'normalization', 'var'), ('batch_stats', 'encoder', 'stages', '3', 'layers', '1', 'layer', 'layer_1', 'normalization', 'mean'), ('params', 'encoder', 'stages', '0', 'layers', '1', 'layer', 'layer_1', 'normalization', 'bias'), ('batch_stats', 'encoder', 'stages', '1', 'layers', '0', 'layer', 'layer_1', 'normalization', 'mean'), ('batch_stats', 'encoder', 'stages', '0', 'layers', '0', 'layer', 'layer_1', 'normalization', 'var'), ('params', 'encoder', 'stages', '2', 'layers', '1', 'layer', 'layer_1', 'normalization', 'scale'), ('params', 'encoder', 'stages', '2', 'layers', '0', 'layer', 'layer_1', 'convolution', 'kernel'), ('params', 'encoder', 'stages', '1', 'layers', '0', 'layer', 'layer_0', 'normalization', 'bias'), ('batch_stats', 'encoder', 'stages', '3', 'layers', '0', 'layer', 'layer_0', 'normalization', 'var'), ('batch_stats', 'encoder', 'stages', '2', 'layers', '1', 'layer', 'layer_0', 'normalization', 'var'), ('params', 'encoder', 'stages', '3', 'layers', '0', 'layer', 'layer_1', 'normalization', 'scale'), ('params', 'encoder', 'stages', '1', 'layers', '1', 'layer', 'layer_1', 'convolution', 'kernel'), ('batch_stats', 'encoder', 'stages', '2', 'layers', '0', 'layer', 'layer_1', 'normalization', 'mean'), ('batch_stats', 'encoder', 'stages', '0', 'layers', '1', 'layer', 'layer_1', 'normalization', 'mean'), ('params', 'encoder', 'stages', '2', 'layers', '0', 'layer', 'layer_0', 'normalization', 'bias'), ('batch_stats', 'encoder', 'stages', '1', 'layers', '1', 'layer', 'layer_0', 'normalization', 'mean'), ('batch_stats', 'encoder', 'stages', '2', 'layers', '1', 'layer', 'layer_1', 'normalization', 'var'), ('batch_stats', 'encoder', 'stages', '0', 'layers', '0', 'layer', 'layer_0', 'normalization', 'var'), ('params', 'encoder', 'stages', '1', 'layers', '0', 'layer', 'layer_1', 'normalization', 'bias'), ('params', 'encoder', 'stages', '1', 'layers', '1', 'layer', 'layer_0', 'normalization', 'bias'), ('params', 'encoder', 'stages', '2', 'layers', '1', 'layer', 'layer_0', 'normalization', 'bias'), ('params', 'encoder', 'stages', '3', 'layers', '0', 'layer', 'layer_0', 'normalization', 'bias'), ('params', 'encoder', 'stages', '2', 'layers', '0', 'layer', 'layer_1', 'normalization', 'bias'), ('batch_stats', 'encoder', 'stages', '1', 'layers', '0', 'layer', 'layer_0', 'normalization', 'mean'), ('params', 'encoder', 'stages', '0', 'layers', '0', 'layer', 'layer_0', 'normalization', 'bias'), ('batch_stats', 'encoder', 'stages', '3', 'layers', '0', 'layer', 'layer_1', 'normalization', 'mean'), ('params', 'encoder', 'stages', '2', 'layers', '1', 'layer', 'layer_1', 'normalization', 'bias'), ('params', 'encoder', 'stages', '1', 'layers', '1', 'layer', 'layer_1', 'normalization', 'scale'), ('params', 'encoder', 'stages', '1', 'layers', '0', 'layer', 'layer_0', 'convolution', 'kernel'), ('params', 'encoder', 'stages', '3', 'layers', '0', 'layer', 'layer_1', 'normalization', 'bias'), ('params', 'encoder', 'stages', '3', 'layers', '1', 'layer', 'layer_0', 'normalization', 'bias'), ('batch_stats', 'encoder', 'stages', '2', 'layers', '0', 'layer', 'layer_0', 'normalization', 'mean'), ('params', 'encoder', 'stages', '0', 'layers', '0', 'layer', 'layer_1', 'normalization', 'bias'), ('batch_stats', 'encoder', 'stages', '0', 'layers', '0', 'layer', 'layer_1', 'normalization', 'mean'), ('batch_stats', 'encoder', 'stages', '2', 'layers', '1', 'layer', 'layer_0', 'normalization', 'mean'), ('batch_stats', 'encoder', 'stages', '3', 'layers', '0', 'layer', 'layer_0', 'normalization', 'mean'), ('params', 'encoder', 'stages', '2', 'layers', '0', 'layer', 'layer_0', 'convolution', 'kernel'), ('batch_stats', 'encoder', 'stages', '0', 'layers', '1', 'layer', 'layer_0', 'normalization', 'var'), ('params', 'encoder', 'stages', '0', 'layers', '1', 'layer', 'layer_0', 'convolution', 'kernel'), ('params', 'encoder', 'stages', '0', 'layers', '0', 'layer', 'layer_1', 'convolution', 'kernel'), ('params', 'encoder', 'stages', '3', 'layers', '1', 'layer', 'layer_1', 'normalization', 'scale'), ('params', 'encoder', 'stages', '2', 'layers', '1', 'layer', 'layer_0', 'convolution', 'kernel'), ('params', 'encoder', 'stages', '3', 'layers', '0', 'layer', 'layer_0', 'convolution', 'kernel'), ('batch_stats', 'encoder', 'stages', '2', 'layers', '1', 'layer', 'layer_1', 'normalization', 'mean'), ('batch_stats', 'encoder', 'stages', '0', 'layers', '0', 'layer', 'layer_0', 'normalization', 'mean'), ('batch_stats', 'encoder', 'stages', '1', 'layers', '1', 'layer', 'layer_1', 'normalization', 'mean'), ('batch_stats', 'encoder', 'stages', '0', 'layers', '1', 'layer', 'layer_1', 'normalization', 'var'), ('params', 'encoder', 'stages', '0', 'layers', '0', 'layer', 'layer_0', 'convolution', 'kernel'), ('params', 'encoder', 'stages', '1', 'layers', '1', 'layer', 'layer_1', 'normalization', 'bias'), ('params', 'encoder', 'stages', '0', 'layers', '1', 'layer', 'layer_1', 'convolution', 'kernel'), ('params', 'encoder', 'stages', '2', 'layers', '1', 'layer', 'layer_1', 'convolution', 'kernel'), ('batch_stats', 'encoder', 'stages', '3', 'layers', '1', 'layer', 'layer_0', 'normalization', 'var'), ('params', 'encoder', 'stages', '0', 'layers', '1', 'layer', 'layer_0', 'normalization', 'scale'), ('batch_stats', 'encoder', 'stages', '1', 'layers', '0', 'layer', 'layer_1', 'normalization', 'var'), ('batch_stats', 'encoder', 'stages', '3', 'layers', '1', 'layer', 'layer_1', 'normalization', 'var'), ('params', 'encoder', 'stages', '3', 'layers', '1', 'layer', 'layer_1', 'normalization', 'bias'), ('params', 'encoder', 'stages', '0', 'layers', '0', 'layer', 'layer_0', 'normalization', 'scale'), ('batch_stats', 'encoder', 'stages', '2', 'layers', '0', 'layer', 'layer_1', 'normalization', 'var'), ('batch_stats', 'encoder', 'stages', '0', 'layers', '1', 'layer', 'layer_0', 'normalization', 'mean'), ('params', 'encoder', 'stages', '0', 'layers', '1', 'layer', 'layer_1', 'normalization', 'scale'), ('batch_stats', 'encoder', 'stages', '1', 'layers', '1', 'layer', 'layer_0', 'normalization', 'var'), ('params', 'encoder', 'stages', '1', 'layers', '0', 'layer', 'layer_0', 'normalization', 'scale'), ('params', 'encoder', 'stages', '3', 'layers', '1', 'layer', 'layer_0', 'normalization', 'scale'), ('params', 'encoder', 'stages', '0', 'layers', '0', 'layer', 'layer_1', 'normalization', 'scale'), ('params', 'encoder', 'stages', '0', 'layers', '1', 'layer', 'layer_0', 'normalization', 'bias'), ('params', 'encoder', 'stages', '2', 'layers', '0', 'layer', 'layer_0', 'normalization', 'scale'), ('params', 'encoder', 'stages', '3', 'layers', '0', 'layer', 'layer_1', 'convolution', 'kernel'), ('batch_stats', 'encoder', 'stages', '1', 'layers', '1', 'layer', 'layer_1', 'normalization', 'var'), ('params', 'encoder', 'stages', '1', 'layers', '0', 'layer', 'layer_1', 'normalization', 'scale'), ('params', 'encoder', 'stages', '3', 'layers', '1', 'layer', 'layer_0', 'convolution', 'kernel'), ('batch_stats', 'encoder', 'stages', '1', 'layers', '0', 'layer', 'layer_0', 'normalization', 'var'), ('params', 'encoder', 'stages', '1', 'layers', '1', 'layer', 'layer_0', 'normalization', 'scale'), ('batch_stats', 'encoder', 'stages', '3', 'layers', '1', 'layer', 'layer_0', 'normalization', 'mean'), ('batch_stats', 'encoder', 'stages', '3', 'layers', '0', 'layer', 'layer_1', 'normalization', 'var'), ('params', 'encoder', 'stages', '2', 'layers', '1', 'layer', 'layer_0', 'normalization', 'scale'), ('params', 'encoder', 'stages', '1', 'layers', '0', 'layer', 'layer_1', 'convolution', 'kernel'), ('params', 'encoder', 'stages', '3', 'layers', '0', 'layer', 'layer_0', 'normalization', 'scale'), ('params', 'encoder', 'stages', '2', 'layers', '0', 'layer', 'layer_1', 'normalization', 'scale'), ('params', 'encoder', 'stages', '1', 'layers', '1', 'layer', 'layer_0', 'convolution', 'kernel'), ('params', 'encoder', 'stages', '3', 'layers', '1', 'layer', 'layer_1', 'convolution', 'kernel')}`


### Expected behavior

I'm expecting to get a similar message to the message i get when loading resnet-26

`Some weights of the model checkpoint at microsoft/resnet-26 were not used when initializing FlaxResNetModel: {('params', 'classifier', '1', 'kernel'), ('params', 'classifier', '1', 'bias')}`

Comments:
We're unfortunately deprecating our Flax code and will be removing it soon, so this will probably not be fixed!

Hi, thanks for the reply. Would you happen to have any suggestions on where I can look for this discrepancy?