Bug: jaxutils.Optimizer.PARAM_COUNTS parameters counts is None about dreamerv3 HOT 2 CLOSED

absl-py==2.1.0
appdirs==1.4.4
astunparse==1.6.3
certifi==2024.2.2
charset-normalizer==3.3.2
chex==0.1.85
click==8.1.7
cloudpickle==1.6.0
crafter==1.8.3
decorator==5.1.1
dm-tree==0.1.8
docker-pycreds==0.4.0
flatbuffers==24.3.7
gast==0.5.4
gitdb==4.0.11
GitPython==3.1.42
google-pasta==0.2.0
grpcio==1.62.1
gym==0.19.0
h5py==3.10.0
idna==3.6
imageio==2.34.0
importlib_metadata==7.0.2
jax==0.4.25
jaxlib==0.4.25+cuda12.cudnn89
keras==3.0.5
libclang==16.0.6
Markdown==3.5.2
markdown-it-py==3.0.0
MarkupSafe==2.1.5
mdurl==0.1.2
ml-dtypes==0.3.2
namex==0.0.7
numpy==1.26.4
nvidia-cublas-cu12==12.4.2.65
nvidia-cuda-cupti-cu12==12.4.99
nvidia-cuda-nvcc-cu12==12.4.99
nvidia-cuda-nvrtc-cu12==12.4.99
nvidia-cuda-runtime-cu12==12.4.99
nvidia-cudnn-cu12==8.9.7.29
nvidia-cufft-cu12==11.2.0.44
nvidia-cusolver-cu12==11.6.0.99
nvidia-cusparse-cu12==12.3.0.142
nvidia-nccl-cu12==2.20.5
nvidia-nvjitlink-cu12==12.4.99
opensimplex==0.4.5
opt-einsum==3.3.0
optax==0.2.1
packaging==24.0
pillow==10.2.0
protobuf==4.25.3
psutil==5.9.8
Pygments==2.17.2
PyYAML==6.0.1
requests==2.31.0
rich==13.7.1
ruamel.yaml==0.17.21
ruamel.yaml.clib==0.2.8
scipy==1.12.0
sentry-sdk==1.42.0
setproctitle==1.3.3
six==1.16.0
smmap==5.0.1
tensorboard==2.16.2
tensorboard-data-server==0.7.2
tensorflow-cpu==2.16.1
tensorflow-io-gcs-filesystem==0.36.0
tensorflow-probability==0.24.0
termcolor==2.4.0
toolz==0.12.1
typing_extensions==4.10.0
urllib3==2.2.1
wandb==0.16.4
Werkzeug==3.0.1
wrapt==1.16.0
zipp==3.18.0