Adapt flax to use any mutable states used by a model generically

[adamgayoso]

No description provided.

[codecov]

Codecov Report

Base: 91.33% // Head: 91.31% // Decreases project coverage by -0.02% ⚠️

Coverage data is based on head (322713e) compared to base (7e45f43).
Patch coverage: 100.00% of modified lines in pull request are covered.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1665      +/-   ##
==========================================
- Coverage   91.33%   91.31%   -0.03%     
==========================================
  Files         116      116              
  Lines        9203     9202       -1     
==========================================
- Hits         8406     8403       -3     
- Misses        797      799       +2     

Impacted Files	Coverage Δ
scvi/module/base/__init__.py	100.00% <100.00%> (ø)
scvi/module/base/_jax_module_wrapper.py	94.05% <100.00%> (ø)
scvi/train/_callbacks.py	83.90% <100.00%> (ø)
scvi/train/_trainingplans.py	95.66% <100.00%> (ø)
scvi/__init__.py	86.66% <0.00%> (-13.34%)	⬇️
scvi/data/_download.py	86.84% <0.00%> (-0.34%)	⬇️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.