A PyTorch library for two-sample tests
Switch branches/tags
Nothing to show
Clone or download
Fetching latest commit…
Cannot retrieve the latest commit at this time.
Failed to load latest commit information.
notebooks Add all requirements of the example in the conda .environment file Mar 7, 2018
.gitignore Use conda environment file Mar 7, 2018
readme.md add travis integration Sep 26, 2017
requirements_docs.txt Add back requirements file for docs and remove minor versions in conda Mar 18, 2018



Documentation Status Build Status

A PyTorch library for differentiable two-sample tests


This package implements a total of six two sample tests:

  • The classical Friedman-Rafsky test [FR79].
  • The classical k-nearest neighbours (kNN) test [FR83].
  • The differentiable Friedman-Rafsky test [DK17].
  • The differentiable k-nearest neighbours (kNN) test [DK17].
  • The maximum mean discrepancy (MMD) test [GBR+12].
  • The energy test [SzekelyR13].

Please refer to the documentation for more information about the project. You can also have a look at the following notebook that showcases how to use the code to train a generative model on MNIST.


After installing PyTorch, you can install the package with:

python setup.py install


To run the tests you simply have to run:

python setup.py test

Note that you will need to have Shogun installed for one of the test cases.


  • [DK17] J. Djolonga and A. Krause. Learning Implicit Generative Models Using Differentiable Graph Tests. ArXiv e-prints, September 2017. arXiv:1709.01006.
  • [FR79] Jerome H Friedman and Lawrence C Rafsky. Multivariate generalizations of the wald-wolfowitz and smirnov two-sample tests. Annals of Statistics, pages 697–717, 1979.
  • [FR83] Jerome H Friedman and Lawrence C Rafsky. Graph-theoretic measures of multivariate association and prediction. Annals of Statistics, pages 377–391, 1983.
  • [GBR+12] Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(Mar):723–773, 2012.
  • [SST+12] Kevin Swersky, Ilya Sutskever, Daniel Tarlow, Richard S Zemel, Ruslan R Salakhutdinov, and Ryan P Adams. Cardinality restricted boltzmann machines. In Advances in Neural Information Processing Systems (NIPS), 3293–3301. 2012.
  • [SzekelyR13] Gábor J Székely and Maria L Rizzo. Energy statistics: a class of statistics based on distances. Journal of Statistical Planning and Inference, 143(8):1249–1272, 2013.
  • [TSZ+12] Daniel Tarlow, Kevin Swersky, Richard S Zemel, Ryan Prescott Adams, and Brendan J Frey. Fast exact inference for recursive cardinality models. Uncertainty in Artificial Intelligence (UAI), 2012.