Allow custom optimizers, improve JaxTrainingPlan

docs/release_notes/index.md

@@ -3,6 +3,14 @@
 This is the list of changes to scvi-tools between each release. Full commit history
 is available in the [commit logs](https://github.com/YosefLab/scvi-tools/commits/).
 
+## Version 0.19
+
+```{toctree}
+:maxdepth: 2
+
+v0.19.0
+```
+
 ## Version 0.18
 
 ```{toctree}

docs/release_notes/v0.19.0.md

@@ -0,0 +1,35 @@
+# New in 0.19.0 (2022-MM-DD)
+
+## Major Changes
+
+-   {class}`~scvi.train.TrainingPlan` allows custom PyTorch optimizers [#1747].
+-   Improvements to {class}`~scvi.train.JaxTrainingPlan` [#1747] [#1749].
+-   {class}`~scvi.module.base.LossRecorder` is deprecated. Please substitute with {class}`~scvi.module.base.LossOutput` [#1749]
+-   All training plans require keyword args after the first positional argument [#1749]
+
+## Minor changes
+
+## Breaking changes
+
+-   {class}`~scvi.module.base.LossRecorder` no longer allows access to dictionaries of values if provided during initialization [#1749].
+
+## Bug Fixes
+
+-   Fix `n_proteins` usage in {class}`~scvi.model.MULTIVI` [#1737].
+-   Remove unused param in {class}`~scvi.model.MULTIVI` [#1741].
+
+## Contributors
+
+-   [@watiss]
+-   [@adamgayoso]
+-   [@martinkim0]
+-   [@marianogabitto]
+
+[#1737]: https://github.com/YosefLab/scvi-tools/pull/1737
+[#1741]: https://github.com/YosefLab/scvi-tools/pull/1737
+[#1747]: https://github.com/YosefLab/scvi-tools/pull/1747
+[#1749]: https://github.com/YosefLab/scvi-tools/pull/1749
+[@watiss]: https://github.com/watiss
+[@adamgayoso]: https://github.com/adamgayoso
+[@martinkim0]: https://github.com/martinkim0
+[@marianogabitto]: https://github.com/marianogabitto

scvi/model/base/_jaxmixin.py

@@ -21,7 +21,7 @@ def train(
         train_size: float = 0.9,
         validation_size: Optional[float] = None,
         batch_size: int = 128,
-        lr: float = 1e-3,
+        plan_kwargs: Optional[dict] = None,
         **trainer_kwargs,
     ):
         """
@@ -43,6 +43,9 @@ def train(
             Minibatch size to use during training.
         lr
             Learning rate to use during training.
+        plan_kwargs
+            Keyword args for :class:`~scvi.train.JaxTrainingPlan`. Keyword arguments passed to
+            `train()` will overwrite values present in `plan_kwargs`, when appropriate.
         **trainer_kwargs
             Other keyword args for :class:`~scvi.train.Trainer`.
         """
@@ -73,10 +76,9 @@ def train(
             use_gpu=False,
             iter_ndarray=True,
         )
+        plan_kwargs = plan_kwargs if isinstance(plan_kwargs, dict) else dict()
 
-        self.training_plan = JaxTrainingPlan(
-            self.module, optim_kwargs=dict(learning_rate=lr)
-        )
+        self.training_plan = JaxTrainingPlan(self.module, **plan_kwargs)
         if "callbacks" not in trainer_kwargs.keys():
             trainer_kwargs["callbacks"] = []
         trainer_kwargs["callbacks"].append(JaxModuleInit())