Add maximum mean discrepancy and radial basis

[kevinchern]

TODOs:

• document
• test
• release note
• wrap with store_config

[kevinchern]

@VolodyaCO what's the motivation for the following changes?

1. Removal of MMD as a module, and
2. The addition of gradient-tracking in the get_bandwidth function.

[kevinchern]

@VolodyaCO what's the motivation for the following changes?

1. Removal of MMD as a module, and

2. The addition of gradient-tracking in the `get_bandwidth` function.

Addressed in meeting. For posterity:

• module was removed for consistency with pseudo_kl_divergence as functions.
• get_bandwidth was accidentally removed (but the l2 matrix is detached in kernel function)

I will review ASAP

[kevinchern]

Thanks Vlad!! Nicely implemented and documented.
Let's add the MaximumMeanDiscrepancy a a module and this should be good to go.

[thisac]

Thanks @VolodyaCO and @kevinchern! Test fail due to Python 3.9 tests being run (removing 3.9 support in #49, which should fix it).

Unit tests should be expanded. There should at least be test classes for kernels (RBF) and more unit tests for the mmd function. Otherwise, looks good. Just a few, mostly minor, comments.

[thisac]

Should Kernels (and RBF) be in a kernels.py instead of mmd.py?

[VolodyaCO]

For now, all kernels are for computing MMD losses. I don't know if that'll change in the future.

[thisac]

If we expect it might change, it seems better to put kernels somewhere else. Alternatively remove Kernel and RBFKernel from __all__ since they're only (?) used within this module.

It seems to me like it's a more general concept and thus, if ever used outside of the mmd module, should be e.g., in torch.model.kernels or torch.kernels even if they're currently only used for calculating MMD losses. Even just putting them in a kernels.py, separating them from the mmd_loss() function, makes more sense to me.

[kevinchern]

Good point Theo. I am in favour of organizing them in torch.kernels for reasons you mentioned, i.e., a kernel is a more general object that isn't limited to applications in MMD.

[VolodyaCO]

In this case, I will move kernels to torch.kernels. Should we have torch.kernels and torch.functional.kernels? To store the function and the modules separately?

[kevinchern]

I can add this^

[thisac]

Do we absolutely need repeat here? Seems a bit cumbersome to add einops as a test dependency just for this test. 🤔

[VolodyaCO]

No, but it's more readable than the pytorch-only version, which requires unsqueezing, inferring the number of feature dimensions and then repeating.

[kevinchern]

chiming in to say einops is generally useful, e.g., #39

[kevinchern]

Just did another pass. Main changes are:

1. Refactor MMD into kernels, functional, and loss modules.
2. Rename acronyms to be consistent with package standards (I'm torn between adhering to our no-acronym standard versus conventions used in pytorch, i.e., pytorch uses MSELoss, but we define MaximuMeanDiscrepancyLoss)
3. Fixed bug with RBF (need to flatten before torch.cdist)
4. Wrote unit tests
5. _kernel(x) changed to _kernel(x, y).
6. Added specific errors: class SampleSizeError(ValueError): ... and class DimensionMismatchError(ValueError): ...

Food for thought (will add an issue): the Kernel base class does not enforce that a kernel is PSD.

Edit RE PSD guarantee: Had a brief exchange with Vlad and figured the onus is on developer to correctly define a kernel

[VolodyaCO]

Looks good to merge

[thisac]

Overall, looks good! Just a few comments/suggestions.