From 76d6fd270dcac3de58fefae916cf257e0d226abe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Paul-Christian=20B=C3=BCrkner?= Date: Thu, 10 Oct 2024 09:22:44 +0200 Subject: [PATCH 1/3] implement student_t.sample natively in keras --- bayesflow/distributions/diagonal_student_t.py | 25 +++++++++++++------ 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/bayesflow/distributions/diagonal_student_t.py b/bayesflow/distributions/diagonal_student_t.py index 8c6a3a7ee..7dc484da3 100644 --- a/bayesflow/distributions/diagonal_student_t.py +++ b/bayesflow/distributions/diagonal_student_t.py @@ -4,7 +4,6 @@ import math import numpy as np -from scipy.stats import t as scipy_student_t from bayesflow.types import Shape, Tensor from .distribution import Distribution @@ -20,6 +19,7 @@ def __init__( loc: int | float | np.ndarray | Tensor = 0.0, scale: int | float | np.ndarray | Tensor = 1.0, use_learnable_parameters: bool = False, + seed_generator: keras.random.SeedGenerator = None, **kwargs, ): super().__init__(**kwargs) @@ -33,6 +33,11 @@ def __init__( self.use_learnable_parameters = use_learnable_parameters + if seed_generator is None: + seed_generator = keras.random.SeedGenerator() + + self.seed_generator = seed_generator + def build(self, input_shape: Shape) -> None: self.dim = int(input_shape[-1]) @@ -78,9 +83,15 @@ def log_prob(self, samples: Tensor, *, normalize: bool = True) -> Tensor: return result def sample(self, batch_shape: Shape) -> Tensor: - # TODO: use reparameterization trick instead of scipy - # TODO: use the seed generator state - dist = scipy_student_t(df=self.df, loc=self.loc, scale=self.scale) - samples = dist.rvs(size=batch_shape + (self.dim,)) - - return keras.ops.convert_to_tensor(samples) + # As of writing this code, keras does not support the chi-square distribution + # nor does it support a scale or rate parameter in Gamma. Hence we use the relation: + # chi-square(df) = Gamma(shape = 0.5 * df, scale = 2) = Gamma(shape = 0.5 * df, scale = 1) * 2 + samples_chisq = keras.random.gamma(batch_shape, alpha=0.5 * self.df) * 2.0 + # the chi-quare samples needss to be repeated across self.dim + # since for each element of batch_shape only one sample is created + samples_chisq = keras.ops.repeat(samples_chisq, self.dim, axis=-1) + samples_chisq = keras.ops.reshape(samples_chisq, batch_shape + (self.dim,)) + + return self.loc + self.scale * keras.random.normal( + shape=batch_shape + (self.dim,), seed=self.seed_generator + ) * keras.ops.sqrt(self.df / samples_chisq) From d5a1aa537d4fa07e87b3db41e9ee15c5a579f5c9 Mon Sep 17 00:00:00 2001 From: larskue Date: Thu, 10 Oct 2024 09:52:07 +0200 Subject: [PATCH 2/3] Fix seed for chi2 samples also do some minor clean-up --- bayesflow/distributions/diagonal_student_t.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/bayesflow/distributions/diagonal_student_t.py b/bayesflow/distributions/diagonal_student_t.py index 7dc484da3..372729f5f 100644 --- a/bayesflow/distributions/diagonal_student_t.py +++ b/bayesflow/distributions/diagonal_student_t.py @@ -6,6 +6,8 @@ from bayesflow.types import Shape, Tensor +from bayesflow.utils import expand_tile + from .distribution import Distribution @@ -84,14 +86,14 @@ def log_prob(self, samples: Tensor, *, normalize: bool = True) -> Tensor: def sample(self, batch_shape: Shape) -> Tensor: # As of writing this code, keras does not support the chi-square distribution - # nor does it support a scale or rate parameter in Gamma. Hence we use the relation: + # nor does it support a scale or rate parameter in Gamma. Hence, we use the relation: # chi-square(df) = Gamma(shape = 0.5 * df, scale = 2) = Gamma(shape = 0.5 * df, scale = 1) * 2 - samples_chisq = keras.random.gamma(batch_shape, alpha=0.5 * self.df) * 2.0 - # the chi-quare samples needss to be repeated across self.dim + chi2_samples = keras.random.gamma(batch_shape, alpha=0.5 * self.df, seed=self.seed_generator) * 2.0 + + # the chi-quare samples needs to be repeated across self.dim # since for each element of batch_shape only one sample is created - samples_chisq = keras.ops.repeat(samples_chisq, self.dim, axis=-1) - samples_chisq = keras.ops.reshape(samples_chisq, batch_shape + (self.dim,)) + chi2_samples = expand_tile(chi2_samples, n=self.dim, axis=-1) + + normal_samples = keras.random.normal(batch_shape + (self.dim,), seed=self.seed_generator) - return self.loc + self.scale * keras.random.normal( - shape=batch_shape + (self.dim,), seed=self.seed_generator - ) * keras.ops.sqrt(self.df / samples_chisq) + return self.loc + self.scale * normal_samples * keras.ops.sqrt(self.df / chi2_samples) From 715c553eba9d0fdf3129c66b30f2c59b2c278abe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Paul-Christian=20B=C3=BCrkner?= Date: Thu, 10 Oct 2024 10:02:29 +0200 Subject: [PATCH 3/3] minor cleaning of comments --- bayesflow/distributions/diagonal_student_t.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bayesflow/distributions/diagonal_student_t.py b/bayesflow/distributions/diagonal_student_t.py index 372729f5f..a02798c4d 100644 --- a/bayesflow/distributions/diagonal_student_t.py +++ b/bayesflow/distributions/diagonal_student_t.py @@ -90,8 +90,8 @@ def sample(self, batch_shape: Shape) -> Tensor: # chi-square(df) = Gamma(shape = 0.5 * df, scale = 2) = Gamma(shape = 0.5 * df, scale = 1) * 2 chi2_samples = keras.random.gamma(batch_shape, alpha=0.5 * self.df, seed=self.seed_generator) * 2.0 - # the chi-quare samples needs to be repeated across self.dim - # since for each element of batch_shape only one sample is created + # The chi-quare samples need to be repeated across self.dim + # since for each element of batch_shape only one sample is created. chi2_samples = expand_tile(chi2_samples, n=self.dim, axis=-1) normal_samples = keras.random.normal(batch_shape + (self.dim,), seed=self.seed_generator)