Migrate flax frozen dicts to dict

[martinkim0]

• Closes Jax tests failing due to module init #2220 (Replace xxxx with the GitHub issue number)
• Tests added and passed if fixing a bug or adding a new feature
• All code checks passed.
• Added type annotations to new arguments/methods/functions.
• Added an entry in the latest docs/release_notes/index.md file if fixing a bug or adding a new feature.
• If the changes are patches for a version, I have added the on-merge: backport to x.x.x label.

[codecov]

Codecov Report

Patch coverage: 100.00% and no project coverage change.

Comparison is base (ac1743b) 89.10% compared to head (6237f46) 89.10%.
Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2222   +/-   ##
=======================================
  Coverage   89.10%   89.10%           
=======================================
  Files         145      145           
  Lines       11809    11810    +1     
=======================================
+ Hits        10522    10523    +1     
  Misses       1287     1287           

Files Changed	Coverage Δ
scvi/external/tangram/_model.py	91.74% <100.00%> (+0.07%)	⬆️
scvi/module/base/_base_module.py	94.25% <100.00%> (-0.03%)	⬇️
scvi/train/_callbacks.py	87.71% <100.00%> (+0.10%)	⬆️

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.