def __init__(self, df, scale, validate_args=None):
self._dim = scale.shape[-1]
assert df > self._dim - 1
self.df = df
self.cholesky_factor = transforms.LowerCholeskyTransform()(scale)
self.chi_sqd_dists = [dist.Chi2(self.df - i) for i in range(self._dim)]
batch_shape, event_shape = scale.shape[:-2], scale.shape[-2:]
super().__init__(batch_shape, event_shape, validate_args)
def __init__(self, multiplicity, in_features, dropout=0.0):
"""Creat a chi square layer.
Args:
multiplicity: Number of parallel representations for each input feature.
in_features: Number of input features.
"""
super().__init__(multiplicity, in_features, dropout)
self.df = nn.Parameter(torch.rand(1, in_features, multiplicity))
self.chi2 = dist.Chi2(df=self.df)
def __init__(self, in_features: int, out_channels: int, num_repetitions: int = 1, dropout=0.0):
"""Creat a chi square layer.
Args:
out_channels: Number of parallel representations for each input feature.
in_features: Number of input features.
num_repetitions: Number of parallel repetitions of this layer.
"""
super().__init__(in_features, out_channels, num_repetitions, dropout)
self.df = nn.Parameter(torch.rand(1, in_features, out_channels, num_repetitions))
self.chi2 = dist.Chi2(df=self.df)
def draw_chisquare(size):
chi_distr = distr.Chi2(1.0)
samples = chi_distr.sample((size,))
return samples
def sample(self, sample_shape=torch.Size()):
# # Step 1: get spherical sample u
# zeros = torch.zeros_like(self.loc)
# ones = torch.ones_like(self.loc)
# sphere_samp = db.Normal(zeros, ones).sample(sample_shape)
# norm_r = (sphere_samp ** 2.0).sum(-1, keepdim=True).sqrt()
# sphere_samp = sphere_samp / norm_r
# local_loc, local_scale, sphere_samp = torch.broadcast_tensors(
# self.loc, self.scale, sphere_samp
# )
# local_cov = self.scale_to_cov(local_scale)
# scale_mat = torch.cholesky(local_cov)
# # Step 2: sample radius
# batch_shape = self.loc.shape[:-1]
# t2_dist = db.Chi2(
# df=torch.tensor(self.d, device=self.loc.device).expand(batch_shape)
# )
# t_samp = t2_dist.sample(sample_shape)
# t_samp = t_samp.sqrt()
# s2_dist = db.Chi2(df=self.df.expand(batch_shape))
# s_samp = s2_dist.rsample(sample_shape)
# s_samp = s_samp.sqrt()
# radius = self.df.sqrt() * t_samp / s_samp
# radius = radius.unsqueeze(-1)
# u = torch.matmul(scale_mat, sphere_samp.unsqueeze(-1)).squeeze(-1)
# assert radius.shape[:-1] == u.shape[:-1]
# samp = local_loc + (radius * u)
# if self.verbose:
# print(
# "z range: {0:.2f} / {1:.2f}".format(
# samp.min().item(), samp.max().item(),
# )
# )
# return samp
# Step 1: get spherical sample u
zeros = torch.zeros_like(self.loc)
ones = torch.ones_like(self.loc)
sphere_samp = db.Normal(zeros, ones).sample(sample_shape)
norm_r = (sphere_samp ** 2.0).sum(-1, keepdim=True).sqrt()
sphere_samp = sphere_samp / norm_r
local_loc, local_scale, sphere_samp = torch.broadcast_tensors(
self.loc, self.scale, sphere_samp
)
scale_mat = local_scale
# Step 2: sample radius
batch_shape = list(self.loc.shape)
batch_shape[-1] = 1 # Radius = 1 parameter
t2_dist = db.Chi2(
df=torch.tensor(self.d, device=self.loc.device).expand(batch_shape)
)
t_samp = t2_dist.sample(sample_shape)
t_samp = t_samp.sqrt()
s2_dist = db.Chi2(df=self.df.expand(batch_shape))
s_samp = s2_dist.rsample(sample_shape)
s_samp = s_samp.sqrt()
radius = self.df.sqrt() * t_samp / s_samp
# radius = radius.unsqueeze(-1)
u = scale_mat * sphere_samp
assert radius.shape[:-1] == u.shape[:-1]
samp = local_loc + (radius * u)
if self.verbose:
print(
"z range: {0:.2f} / {1:.2f}".format(
samp.min().item(), samp.max().item(),
)
)
return samp
