Backport PR #1665 on branch 0.17.x (Adapt flax to use any mutable states used by a model generically)

[meeseeksmachine]

Backport PR #1665: Adapt flax to use any mutable states used by a model generically

[codecov]

Codecov Report

Base: 91.33% // Head: 91.33% // No change to project coverage 👍

Coverage data is based on head (49c6045) compared to base (b1398b1).
Patch coverage: 100.00% of modified lines in pull request are covered.

Additional details and impacted files

@@           Coverage Diff           @@
##           0.17.x    #1699   +/-   ##
=======================================
  Coverage   91.33%   91.33%           
=======================================
  Files         116      116           
  Lines        9203     9203           
=======================================
  Hits         8406     8406           
  Misses        797      797           

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%> (ø)

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.