AttributeError: module 'flax' has no attribute 'nn'

Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.

### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: 0.7.3 , 0.4.14 , 0.4.14
- Python version: Python 3.9.12
- GPU/TPU model and memory: NVIDIA GeForce
- CUDA version (if applicable): 11.7

### Problem you have encountered:
im working in mip-nerf form google and tried to train some dataset and become Error:
2023-08-27 13:09:05.370975: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
Traceback (most recent call last):
File "/home/alhasan.ali/mip/train.py", line 31, in
from internal import datasets
File "/home/alhasan.ali/mip/internal/datasets.py", line 26, in
from internal import utils
File "/home/alhasan.ali/mip/internal/utils.py", line 32, in
gin.config.external_configurable(flax.nn.relu, module='flax.nn')
AttributeError: module 'flax' has no attribute 'nn'

### What you expected to happen:

### Logs, error messages, etc:

### Steps to reproduce:
Whenever possible, please provide a *minimal example*. Please consider submitting it as a Colab link.