60 add vi helper

[Tennessee-Wallaceh]

Initial solution to #60.

Key points:

• Creates a new flowjax.train submodule, for organisational purposes. The old train utils are moved to flowjax.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 simpler

• I am currently adding some tests

[danielward27]

Nice, I can add in the sample_and_log_prob method.

[danielward27]

I've added some comments, thanks again for adding this :)

[danielward27]

I've added the sample_and_log_prob method in #65, so the training script should be updated to use that.

[danielward27]

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.

[Tennessee-Wallaceh]

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?