Closed Tennessee-Wallaceh closed 1 year ago
Nice, I can add in the sample_and_log_prob method.
I've added the sample_and_log_prob
method in https://github.com/danielward27/flowjax/pull/65, so the training script should be updated to use that.
By the way, I have no strong feelings about "train" vs "fit", but feel that we should be consistent and use one or the other and not both.
By the way, I have no strong feelings about "train" vs "fit", but feel that we should be consistent and use one or the other and not both.
How about fit_to_data
and fit_to_variational_target
?
Initial solution to #60.
Key points:
Creates a new
flowjax.train
submodule, for organisational purposes. The old train utils are moved toflowjax.train.data_fit
New VI helper is in
flowjax.train.variational_fit
Doesn't include a
sample_and_log_prob
method, which can come in separate work to keep the PRs simplerI am currently adding some tests