def distribution(
self,
distr_args,
loc: Optional[Tensor] = None,
scale: Optional[Tensor] = None,
) -> Distribution:
r"""
Construct the associated distribution, given the collection of
constructor arguments and, optionally, a scale tensor.
Parameters
----------
distr_args
Constructor arguments for the underlying Distribution type.
loc
Optional tensor, of the same shape as the
batch_shape+event_shape of the resulting distribution.
scale
Optional tensor, of the same shape as the
batch_shape+event_shape of the resulting distribution.
"""
if loc is None and scale is None:
return self.distr_cls(*distr_args)
else:
distr = self.distr_cls(*distr_args)
return TransformedDistribution(
distr, [AffineTransformation(loc=loc, scale=scale)])
def distribution(self,
distr_args,
scale: Optional[Tensor] = None,
**kwargs) -> PiecewiseLinear:
if scale is None:
return self.distr_cls(*distr_args)
else:
distr = self.distr_cls(*distr_args)
return TransformedPiecewiseLinear(
distr, AffineTransformation(scale=scale))
def distribution(self,
distr_args,
scale: Optional[Tensor] = None) -> Distribution:
r"""
Construct the associated distribution, given the collection of
constructor arguments and, optionally, a scale tensor.
"""
if scale is None:
return self.distr_cls(*distr_args)
else:
distr = self.distr_cls(*distr_args)
return TransformedDistribution(distr,
AffineTransformation(scale=scale))
def distribution(self,
distr_args,
scale: Optional[Tensor] = None,
**kwargs) -> Distribution:
distr_args, transforms_args = self._split_args(distr_args)
distr = self.base_distr_output.distr_cls(*distr_args)
transforms = [
transform_output.bij_cls(*bij_args) for transform_output, bij_args
in zip(self.transforms_output, transforms_args)
]
trans_distr = TransformedDistribution(distr, transforms)
# Apply scaling as well at the end if scale is not None!
if scale is None:
return trans_distr
else:
return TransformedDistribution(trans_distr,
AffineTransformation(scale=scale))
F=mx.nd, x=mx.nd.ones(shape=(3, 4, 5)), d=5),
),
],
),
(3, 4),
(5, ),
),
(
TransformedDistribution(
StudentT(
mu=mx.nd.zeros(shape=(3, 4, 5)),
sigma=mx.nd.ones(shape=(3, 4, 5)),
nu=mx.nd.ones(shape=(3, 4, 5)),
),
[
AffineTransformation(scale=1e-1 +
mx.nd.random.uniform(shape=(3, 4, 5)))
],
),
(3, 4, 5),
(),
),
(
TransformedDistribution(
MultivariateGaussian(
mu=mx.nd.zeros(shape=(3, 4, 5)),
L=make_nd_diag(F=mx.nd, x=mx.nd.ones(shape=(3, 4, 5)),
d=5),
),
[
AffineTransformation(scale=1e-1 +
mx.nd.random.uniform(shape=(3, 4, 5)))