tfp.substrates.jax.distributions.kl_divergence

Get the KL-divergence KL(distribution_a || distribution_b).

tfp.substrates.jax.distributions.kl_divergence(
    distribution_a, distribution_b, allow_nan_stats=True, name=None
)

If there is no KL method registered specifically for type(distribution_a) and type(distribution_b), then the class hierarchies of these types are searched.

If one KL method is registered between any pairs of classes in these two parent hierarchies, it is used.

If more than one such registered method exists, the method whose registered classes have the shortest sum MRO paths to the input types is used.

If more than one such shortest path exists, the first method identified in the search is used (favoring a shorter MRO distance to type(distribution_a)).

distribution_a The first distribution.
distribution_b The second distribution.
allow_nan_stats Python bool, default True. When True, statistics (e.g., mean, mode, variance) use the value "NaN" to indicate the result is undefined. When False, an exception is raised if one or more of the statistic's batch members are undefined.
name Python str name prefixed to Ops created by this class.

A Tensor with the batchwise KL-divergence between distribution_a and distribution_b.

NotImplementedError If no KL method is defined for distribution types of distribution_a and distribution_b.

Built-in KL(distribution_a || distribution_b) registrations:

                          distribution_a || distribution_b
====================================================================================
                               Bernoulli || Bernoulli
                                    Beta || Beta
                             Categorical || Categorical
                                  Cauchy || Cauchy
                                     Chi || Chi
                                    Chi2 || Chi2
                                    Chi2 || Gamma
                                            + Exponential
                     ContinuousBernoulli || ContinuousBernoulli
                               Dirichlet || Dirichlet
                                ExpGamma || ExpGamma
                         ExpInverseGamma || ExpInverseGamma
                                   Gamma || Chi2
                           Exponential +    
                                   Gamma || Gamma
                           Exponential +    + Exponential
                         GaussianProcess || MultivariateNormalLinearOperator
        GaussianProcessRegressionModel +    + MultivariateNormalDiag
            VariationalGaussianProcess +    + MultivariateNormalDiagPlusLowRank
                                            + MultivariateNormalFullCovariance
                                            + MultivariateNormalTriL
                                  Gumbel || Gumbel
                              HalfNormal || HalfNormal
                                 Laplace || Laplace
                               LogNormal || LogNormal
                             LogitNormal || LogitNormal
              MatrixNormalLinearOperator || MatrixNormalLinearOperator
                                   Moyal || Moyal
        MultivariateNormalLinearOperator || GaussianProcess
                MultivariateNormalDiag +    + GaussianProcessRegressionModel
     MultivariateNormalDiagPlusLowRank +    + VariationalGaussianProcess
      MultivariateNormalFullCovariance +    
                MultivariateNormalTriL +    
        MultivariateNormalLinearOperator || MultivariateNormalLinearOperator
                MultivariateNormalDiag +    + MultivariateNormalDiag
     MultivariateNormalDiagPlusLowRank +    + MultivariateNormalDiagPlusLowRank
      MultivariateNormalFullCovariance +    + MultivariateNormalFullCovariance
                MultivariateNormalTriL +    + MultivariateNormalTriL
                                  Normal || Normal
                       OneHotCategorical || OneHotCategorical
                         OrderedLogistic || OrderedLogistic
                                  Pareto || Pareto
                          PowerSpherical || SphericalUniform
                          PowerSpherical || VonMisesFisher
                         ProbitBernoulli || ProbitBernoulli
                   StoppingRatioLogistic || StoppingRatioLogistic
                                 Uniform || Uniform
                                VonMises || VonMises
                          VonMisesFisher || SphericalUniform
                                 Weibull || Gamma
                                            + Exponential
                                 Weibull || Weibull
                      _BaseDeterministic || Distribution
                         Deterministic +    + AutoCompositeTensorDistribution
                   VectorDeterministic +    + Autoregressive
                                            + BatchBroadcast
                                            + BatchConcat
                                            + BatchReshape
                                            + Bates
                                            + 116 more
                              _Blockwise || _Blockwise
                             Blockwise +    + Blockwise
                            _Independent || _Independent
                           Independent +    + Independent
            _JointDistributionSequential || _JointDistributionSequential
                JointDistributionNamed +    + JointDistributionNamed
     JointDistributionNamedAutoBatched +    + JointDistributionNamedAutoBatched
           JointDistributionSequential +    + JointDistributionSequential
JointDistributionSequentialAutoBatched +    + JointDistributionSequentialAutoBatched
                                 _Masked || _Masked
                                Masked +    + Masked
                 _NonCompositeTensorCast || _NonCompositeTensorCast
                                 _Sample || _Sample
                                Sample +    + Sample
                _TransformedDistribution || _TransformedDistribution
                                   Chi +    + Chi
                       ExpInverseGamma +    + ExpInverseGamma
               GeneralizedExtremeValue +    + GeneralizedExtremeValue
                                Gumbel +    + Gumbel
                             JohnsonSU +    + JohnsonSU
                           Kumaraswamy +    + Kumaraswamy
                               15 more +    + 15 more