Remove need for JaxModuleWrapper

[justjhong]

As mentioned in https://github.com/scverse/scvi-tools/issues/1564, if we find a way to default training without explicitly passing it in, this should be enough to move the necessary logic from JaxModuleWrapper to JaxBaseModuleClass and deprecate JaxModuleWrapper entirely. This interface will be much more consistent with BaseModuleClass.