diff --git a/tensorflow_probability/python/distributions/BUILD b/tensorflow_probability/python/distributions/BUILD index c09884696f..0d8c82f760 100644 --- a/tensorflow_probability/python/distributions/BUILD +++ b/tensorflow_probability/python/distributions/BUILD @@ -144,6 +144,7 @@ multi_substrate_py_library( ":sigmoid_beta", ":sinh_arcsinh", ":skellam", + ":skew_normal", ":spherical_uniform", ":stopping_ratio_logistic", ":student_t", @@ -2372,6 +2373,26 @@ multi_substrate_py_library( ], ) +multi_substrate_py_library( + name = "skew_normal", + srcs = ["skew_normal.py"], + deps = [ + ":distribution", + # numpy dep, + # tensorflow dep, + "//tensorflow_probability/python/bijectors:identity", + "//tensorflow_probability/python/bijectors:softplus", + "//tensorflow_probability/python/internal:assert_util", + "//tensorflow_probability/python/internal:dtype_util", + "//tensorflow_probability/python/internal:parameter_properties", + "//tensorflow_probability/python/internal:prefer_static", + "//tensorflow_probability/python/internal:reparameterization", + "//tensorflow_probability/python/internal:samplers", + "//tensorflow_probability/python/internal:tensor_util", + "//tensorflow_probability/python/math:special", + ], +) + multi_substrate_py_library( name = "spherical_uniform", srcs = ["spherical_uniform.py"], @@ -4707,6 +4728,18 @@ multi_substrate_py_test( ], ) +multi_substrate_py_test( + name = "skew_normal_test", + srcs = ["skew_normal_test.py"], + deps = [ + ":skew_normal", + # numpy dep, + # scipy dep, + # tensorflow dep, + "//tensorflow_probability/python/internal:test_util", + ], +) + multi_substrate_py_test( name = "spherical_uniform_test", srcs = ["spherical_uniform_test.py"], diff --git a/tensorflow_probability/python/distributions/__init__.py b/tensorflow_probability/python/distributions/__init__.py index 6160c7d2f6..268ab15ee8 100644 --- a/tensorflow_probability/python/distributions/__init__.py +++ b/tensorflow_probability/python/distributions/__init__.py @@ -122,6 +122,7 @@ from tensorflow_probability.python.distributions.sigmoid_beta import SigmoidBeta from tensorflow_probability.python.distributions.sinh_arcsinh import SinhArcsinh from tensorflow_probability.python.distributions.skellam import Skellam +from tensorflow_probability.python.distributions.skew_normal import SkewNormal from tensorflow_probability.python.distributions.spherical_uniform import SphericalUniform from tensorflow_probability.python.distributions.stopping_ratio_logistic import StoppingRatioLogistic from tensorflow_probability.python.distributions.student_t import StudentT @@ -282,6 +283,7 @@ 'SigmoidBeta', 'SinhArcsinh', 'Skellam', + 'SkewNormal', 'SphericalUniform', 'StoppingRatioLogistic', 'StudentT', diff --git a/tensorflow_probability/python/distributions/skew_normal.py b/tensorflow_probability/python/distributions/skew_normal.py new file mode 100644 index 0000000000..a061c08da0 --- /dev/null +++ b/tensorflow_probability/python/distributions/skew_normal.py @@ -0,0 +1,221 @@ +# Copyright 2026 The TensorFlow Probability Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""The Skew-Normal distribution class.""" + +import numpy as np +import tensorflow.compat.v2 as tf + +from tensorflow_probability.python.bijectors import identity as identity_bijector +from tensorflow_probability.python.bijectors import softplus as softplus_bijector +from tensorflow_probability.python.distributions import distribution +from tensorflow_probability.python.internal import assert_util +from tensorflow_probability.python.internal import dtype_util +from tensorflow_probability.python.internal import parameter_properties +from tensorflow_probability.python.internal import prefer_static as ps +from tensorflow_probability.python.internal import reparameterization +from tensorflow_probability.python.internal import samplers +from tensorflow_probability.python.internal import special_math +from tensorflow_probability.python.internal import tensor_util +from tensorflow_probability.python.math import special as tfp_math_special + +__all__ = [ + 'SkewNormal', +] + +class SkewNormal(distribution.AutoCompositeTensorDistribution): + """The Skew-Normal distribution. + + The Skew-Normal distribution is a generalization of the Normal distribution + that allows for non-zero skewness. + + #### Mathematical details + + The probability density function (pdf) is, + + ```none + pdf(x; loc, scale, skewness) = 2 / scale * phi(z) * Phi(skewness * z) + z = (x - loc) / scale + ``` + + where: + * `loc` is the location parameter, + * `scale` is the scale parameter, + * `skewness` (often denoted as `alpha`) is the skewness parameter, + * `phi(z)` is the standard normal pdf, and + * `Phi(z)` is the standard normal cdf. + + The cumulative distribution function (cdf) is, + + ```none + cdf(x; loc, scale, skewness) = Phi(z) - 2 * T(z, skewness) + ``` + + where `T` is Owen's T function. + """ + + def __init__(self, + loc, + scale, + skewness, + validate_args=False, + allow_nan_stats=True, + name='SkewNormal'): + """Construct Skew-Normal distribution with parameters `loc`, `scale`, and `skewness`. + + Args: + loc: Floating point tensor; the location of the distribution(s). + scale: Floating point tensor; the scale of the distribution(s). + Must contain only positive values. + skewness: Floating point tensor; the skewness of the distribution(s). + validate_args: Python `bool`, default `False`. When `True` distribution + parameters are checked for validity despite possibly degrading runtime + performance. When `False` invalid inputs may silently render incorrect + outputs. + 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. + """ + parameters = dict(locals()) + with tf.name_scope(name) as name: + dtype = dtype_util.common_dtype( + [loc, scale, skewness], dtype_hint=tf.float32) + self._loc = tensor_util.convert_nonref_to_tensor( + loc, dtype=dtype, name='loc') + self._scale = tensor_util.convert_nonref_to_tensor( + scale, dtype=dtype, name='scale') + self._skewness = tensor_util.convert_nonref_to_tensor( + skewness, dtype=dtype, name='skewness') + + super(SkewNormal, self).__init__( + dtype=dtype, + reparameterization_type=reparameterization.FULLY_REPARAMETERIZED, + validate_args=validate_args, + allow_nan_stats=allow_nan_stats, + parameters=parameters, + name=name) + + @classmethod + def _parameter_properties(cls, dtype, num_classes=None): + # pylint: disable=g-long-lambda + return dict( + loc=parameter_properties.ParameterProperties(), + scale=parameter_properties.ParameterProperties( + default_constraining_bijector_fn=( + lambda: softplus_bijector.Softplus(low=dtype_util.eps(dtype)))), + skewness=parameter_properties.ParameterProperties()) + # pylint: enable=g-long-lambda + + @property + def loc(self): + """Distribution parameter for the location.""" + return self._loc + + @property + def scale(self): + """Distribution parameter for scale.""" + return self._scale + + @property + def skewness(self): + """Distribution parameter for skewness.""" + return self._skewness + + def _event_shape_tensor(self): + return tf.constant([], dtype=tf.int32) + + def _event_shape(self): + return tf.TensorShape([]) + + def _sample_n(self, n, seed=None): + loc = tf.convert_to_tensor(self.loc) + scale = tf.convert_to_tensor(self.scale) + skewness = tf.convert_to_tensor(self.skewness) + + shape = ps.concat([[n], self._batch_shape_tensor( + loc=loc, scale=scale, skewness=skewness)], axis=0) + + seed1, seed2 = samplers.split_seed(seed, salt='skew_normal') + + u0 = samplers.normal(shape=shape, mean=0., stddev=1., dtype=self.dtype, seed=seed1) + v = samplers.normal(shape=shape, mean=0., stddev=1., dtype=self.dtype, seed=seed2) + + delta = skewness / tf.math.sqrt(1. + tf.math.square(skewness)) + z = delta * tf.math.abs(u0) + tf.math.sqrt(1. - tf.math.square(delta)) * v + + return z * scale + loc + + def _log_prob(self, x): + scale = tf.convert_to_tensor(self.scale) + skewness = tf.convert_to_tensor(self.skewness) + z = (x - self.loc) / scale + + log_unnormalized = -0.5 * tf.math.square(z) + log_normalization = tf.constant( + 0.5 * np.log(2. * np.pi), dtype=self.dtype) + tf.math.log(scale) + log_pdf_normal = log_unnormalized - log_normalization + + log_cdf_skew = special_math.log_ndtr(skewness * z) + + return tf.constant(np.log(2.), dtype=self.dtype) + log_pdf_normal + log_cdf_skew + + def _cdf(self, x): + scale = tf.convert_to_tensor(self.scale) + skewness = tf.convert_to_tensor(self.skewness) + z = (x - self.loc) / scale + + return special_math.ndtr(z) - 2. * tfp_math_special.owens_t(z, skewness) + + def _mean(self): + scale = tf.convert_to_tensor(self.scale) + skewness = tf.convert_to_tensor(self.skewness) + delta = skewness / tf.math.sqrt(1. + tf.math.square(skewness)) + return self.loc + scale * delta * tf.constant(np.sqrt(2. / np.pi), dtype=self.dtype) + + def _variance(self): + scale = tf.convert_to_tensor(self.scale) + skewness = tf.convert_to_tensor(self.skewness) + delta = skewness / tf.math.sqrt(1. + tf.math.square(skewness)) + + c = tf.constant(2. / np.pi, dtype=self.dtype) + return tf.math.square(scale) * (1. - c * tf.math.square(delta)) + + def _default_event_space_bijector(self): + return identity_bijector.Identity(validate_args=self.validate_args) + + def _parameter_control_dependencies(self, is_init): + assertions = [] + + if is_init: + try: + self._batch_shape() + except ValueError: + raise ValueError( + 'Arguments `loc`, `scale`, and `skewness` must have compatible ' + 'shapes; loc.shape={}, scale.shape={}, skewness.shape={}.'.format( + self.loc.shape, self.scale.shape, self.skewness.shape)) + # We don't bother checking the shapes in the dynamic case because + # all member functions access both arguments anyway. + + if not self.validate_args: + assert not assertions # Should never happen. + return [] + + if is_init != tensor_util.is_ref(self.scale): + assertions.append(assert_util.assert_positive( + self.scale, message='Argument `scale` must be positive.')) + + return assertions diff --git a/tensorflow_probability/python/distributions/skew_normal_test.py b/tensorflow_probability/python/distributions/skew_normal_test.py new file mode 100644 index 0000000000..5eb387eb5f --- /dev/null +++ b/tensorflow_probability/python/distributions/skew_normal_test.py @@ -0,0 +1,138 @@ +# Copyright 2026 The TensorFlow Probability Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""Tests for SkewNormal.""" + +import math +import numpy as np +from scipy import stats as sp_stats + +import tensorflow.compat.v2 as tf + +from tensorflow_probability.python.distributions import skew_normal +from tensorflow_probability.python.internal import test_util + + +@test_util.test_all_tf_execution_regimes +class SkewNormalTest(test_util.TestCase): + + def setUp(self): + self._rng = np.random.RandomState(123) + super(SkewNormalTest, self).setUp() + + def testSkewNormalShape(self): + loc = tf.constant([3.0] * 5) + scale = tf.constant(11.0) + skewness = tf.constant([2.0] * 5) + dist = skew_normal.SkewNormal(loc=loc, scale=scale, skewness=skewness) + + self.assertEqual(self.evaluate(dist.batch_shape_tensor()), (5,)) + self.assertEqual(dist.batch_shape, tf.TensorShape([5])) + self.assertAllEqual(self.evaluate(dist.event_shape_tensor()), []) + self.assertEqual(dist.event_shape, tf.TensorShape([])) + + def testSkewNormalLogPDF(self): + batch_size = 6 + loc = tf.constant([2.0] * batch_size) + scale = tf.constant([3.0] * batch_size) + skewness = tf.constant([0.0, 1.0, -1.0, 2.0, -2.0, 5.0]) + x = np.array([2.5, 2.5, 4.0, -1.0, 5.0, 2.0], dtype=np.float32) + + dist = skew_normal.SkewNormal(loc=loc, scale=scale, skewness=skewness) + log_pdf = dist.log_prob(x) + + expected_log_pdf = sp_stats.skewnorm.logpdf( + x, a=self.evaluate(skewness), loc=self.evaluate(loc), scale=self.evaluate(scale)) + + self.assertAllClose(self.evaluate(log_pdf), expected_log_pdf, rtol=1e-4) + + def testSkewNormalCDF(self): + batch_size = 6 + loc = tf.constant([2.0] * batch_size) + scale = tf.constant([3.0] * batch_size) + skewness = tf.constant([0.0, 1.0, -1.0, 2.0, -2.0, 5.0]) + x = np.array([2.5, 2.5, 4.0, -1.0, 5.0, 2.0], dtype=np.float32) + + dist = skew_normal.SkewNormal(loc=loc, scale=scale, skewness=skewness) + cdf = dist.cdf(x) + + expected_cdf = sp_stats.skewnorm.cdf( + x, a=self.evaluate(skewness), loc=self.evaluate(loc), scale=self.evaluate(scale)) + + self.assertAllClose(self.evaluate(cdf), expected_cdf, rtol=1e-4) + + def testSkewNormalMean(self): + loc = np.array([2.0, -1.0, 0.0]) + scale = np.array([3.0, 0.5, 1.0]) + skewness = np.array([1.0, -2.0, 0.0]) + + dist = skew_normal.SkewNormal(loc=loc, scale=scale, skewness=skewness) + expected_mean = sp_stats.skewnorm.mean(a=skewness, loc=loc, scale=scale) + + self.assertAllClose(self.evaluate(dist.mean()), expected_mean, rtol=1e-4) + + def testSkewNormalVariance(self): + loc = np.array([2.0, -1.0, 0.0]) + scale = np.array([3.0, 0.5, 1.0]) + skewness = np.array([1.0, -2.0, 0.0]) + + dist = skew_normal.SkewNormal(loc=loc, scale=scale, skewness=skewness) + expected_var = sp_stats.skewnorm.var(a=skewness, loc=loc, scale=scale) + + self.assertAllClose(self.evaluate(dist.variance()), expected_var, rtol=1e-4) + + def testSkewNormalSample(self): + loc = tf.constant(2.0) + scale = tf.constant(3.0) + skewness = tf.constant(1.5) + + dist = skew_normal.SkewNormal(loc=loc, scale=scale, skewness=skewness) + + n = 100000 + samples = dist.sample(n, seed=test_util.test_seed()) + sample_values = self.evaluate(samples) + + self.assertEqual(sample_values.shape, (n,)) + self.assertAllClose( + sample_values.mean(), + self.evaluate(dist.mean()), + rtol=0.02, + atol=0.02) + self.assertAllClose( + sample_values.var(), + self.evaluate(dist.variance()), + rtol=0.02, + atol=0.02) + + def testSkewNormalFullyReparameterized(self): + loc = tf.constant(2.0) + scale = tf.constant(3.0) + skewness = tf.constant(1.5) + + # We do a simpler reparameterization check instead of `compute_gradient` + # because tf.test.compute_gradient expects scalar output or specific shapes. + # Instead, we just check if the gradient w.r.t parameters is not None. + with tf.GradientTape() as tape: + tape.watch([loc, scale, skewness]) + dist = skew_normal.SkewNormal(loc=loc, scale=scale, skewness=skewness) + samples = dist.sample(100, seed=test_util.test_seed()) + loss = tf.reduce_mean(samples) + + grad_loc, grad_scale, grad_skewness = tape.gradient(loss, [loc, scale, skewness]) + self.assertIsNotNone(grad_loc) + self.assertIsNotNone(grad_scale) + self.assertIsNotNone(grad_skewness) + +if __name__ == '__main__': + test_util.main()