pixyz.distributions (Distribution API)

Distribution

class pixyz.distributions.distributions.Distribution(var, cond_var=[], name='p', features_shape=torch.Size([]))[source]

Bases: torch.nn.modules.module.Module

Distribution class. In Pixyz, all distributions are required to inherit this class.

Examples

>>> import torch
>>> from torch.nn import functional as F
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[64], name="p1")
>>> print(p1)
Distribution:
  p_{1}(x)
Network architecture:
  Normal(
    name=p_{1}, distribution_name=Normal,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([64])
    (loc): torch.Size([1, 64])
    (scale): torch.Size([1, 64])
  )
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[64], name="p2")
>>> print(p2)
Distribution:
  p_{2}(x|y)
Network architecture:
  Normal(
    name=p_{2}, distribution_name=Normal,
    var=['x'], cond_var=['y'], input_var=['y'], features_shape=torch.Size([64])
    (scale): torch.Size([1, 64])
  )
>>> # Conditional distribution (by neural networks)
>>> class P(Normal):
...     def __init__(self):
...         super().__init__(var=["x"], cond_var=["y"], name="p3")
...         self.model_loc = nn.Linear(128, 64)
...         self.model_scale = nn.Linear(128, 64)
...     def forward(self, y):
...         return {"loc": self.model_loc(y), "scale": F.softplus(self.model_scale(y))}
>>> p3 = P()
>>> print(p3)
Distribution:
  p_{3}(x|y)
Network architecture:
  P(
    name=p_{3}, distribution_name=Normal,
    var=['x'], cond_var=['y'], input_var=['y'], features_shape=torch.Size([])
    (model_loc): Linear(in_features=128, out_features=64, bias=True)
    (model_scale): Linear(in_features=128, out_features=64, bias=True)
  )
__init__(var, cond_var=[], name='p', features_shape=torch.Size([]))[source]
Parameters:
  • var (list of str) – Variables of this distribution.
  • cond_var (list of str, defaults to []) – Conditional variables of this distribution. In case that cond_var is not empty, we must set the corresponding inputs to sample variables.
  • name (str, defaults to “p”) – Name of this distribution. This name is displayed in prob_text and prob_factorized_text.
  • features_shape (torch.Size or list, defaults to torch.Size())) – Shape of dimensions (features) of this distribution.
distribution_name

Name of this distribution class.

Type:str
name

Name of this distribution displayed in prob_text and prob_factorized_text.

Type:str
var

Variables of this distribution.

Type:list
cond_var

Conditional variables of this distribution.

Type:list
input_var

Input variables of this distribution. Normally, it has same values as cond_var.

Type:list
prob_text

Return a formula of the (joint) probability distribution.

Type:str
prob_factorized_text

Return a formula of the factorized probability distribution.

Type:str
prob_joint_factorized_and_text

Return a formula of the factorized and the (joint) probability distributions.

Type:str
features_shape

Shape of features of this distribution.

Type:torch.Size or list
get_params(params_dict={})[source]

This method aims to get parameters of this distributions from constant parameters set in initialization and outputs of DNNs.

Parameters:params_dict (dict, defaults to {}) – Input parameters.
Returns:output_dict – Output parameters.
Return type:dict

Examples

>>> from pixyz.distributions import Normal
>>> dist_1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...                 features_shape=[1])
>>> print(dist_1)
Distribution:
  p(x)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([1])
    (loc): torch.Size([1, 1])
    (scale): torch.Size([1, 1])
  )
>>> dist_1.get_params()
{'loc': tensor([[0.]]), 'scale': tensor([[1.]])}
>>> dist_2 = Normal(loc=torch.tensor(0.), scale="z", cond_var=["z"], var=["x"])
>>> print(dist_2)
Distribution:
  p(x|z)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
    (loc): torch.Size([1])
  )
>>> dist_2.get_params({"z": torch.tensor(1.)})
{'scale': tensor(1.), 'loc': tensor([0.])}
sample(x_dict={}, batch_n=None, sample_shape=torch.Size([]), return_all=True, reparam=False)[source]

Sample variables of this distribution. If cond_var is not empty, you should set inputs as dict.

Parameters:
  • x_dict (torch.Tensor, list, or dict, defaults to {}) – Input variables.
  • sample_shape (list or NoneType, defaults to torch.Size()) – Shape of generating samples.
  • batch_n (int, defaults to None.) – Set batch size of parameters.
  • return_all (bool, defaults to True) – Choose whether the output contains input variables.
  • reparam (bool, defaults to False.) – Choose whether we sample variables with re-parameterized trick.
Returns:

output – Samples of this distribution.

Return type:

dict

Examples

>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...            features_shape=[10, 2])
>>> print(p)
Distribution:
  p(x)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([10, 2])
    (loc): torch.Size([1, 10, 2])
    (scale): torch.Size([1, 10, 2])
  )
>>> p.sample()["x"].shape  # (batch_n=1, features_shape)
torch.Size([1, 10, 2])
>>> p.sample(batch_n=20)["x"].shape  # (batch_n, features_shape)
torch.Size([20, 10, 2])
>>> p.sample(batch_n=20, sample_shape=[40, 30])["x"].shape  # (sample_shape, batch_n, features_shape)
torch.Size([40, 30, 20, 10, 2])
>>> # Conditional distribution
>>> p = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...            features_shape=[10])
>>> print(p)
Distribution:
  p(x|y)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=['y'], input_var=['y'], features_shape=torch.Size([10])
    (scale): torch.Size([1, 10])
  )
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> sample_a = torch.randn(1, 10) # Psuedo data
>>> sample = p.sample({"y": sample_y})
>>> print(sample) # input_var + var  # doctest: +SKIP
{'y': tensor([[-0.5182,  0.3484,  0.9042,  0.1914,  0.6905,
               -1.0859, -0.4433, -0.0255, 0.8198,  0.4571]]),
 'x': tensor([[-0.7205, -1.3996,  0.5528, -0.3059,  0.5384,
               -1.4976, -0.1480,  0.0841,0.3321,  0.5561]])}
>>> sample = p.sample({"y": sample_y, "a": sample_a}) # Redundant input ("a")
>>> print(sample) # input_var + var + "a" (redundant input)  # doctest: +SKIP
{'y': tensor([[ 1.3582, -1.1151, -0.8111,  1.0630,  1.1633,
                0.3855,  2.6324, -0.9357, -0.8649, -0.6015]]),
 'a': tensor([[-0.1874,  1.7958, -1.4084, -2.5646,  1.0868,
               -0.7523, -0.0852, -2.4222, -0.3914, -0.9755]]),
 'x': tensor([[-0.3272, -0.5222, -1.3659,  1.8386,  2.3204,
                0.3686,  0.6311, -1.1208, 0.3656, -0.6683]])}
sample_mean(x_dict={})[source]

Return the mean of the distribution.

Parameters:x_dict (dict, defaults to {}) – Parameters of this distribution.

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> mean = p1.sample_mean()
>>> print(mean)
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> mean = p2.sample_mean({"y": sample_y})
>>> print(mean) # doctest: +SKIP
tensor([[-0.2189, -1.0310, -0.1917, -0.3085,  1.5190, -0.9037,  1.2559,  0.1410,
          1.2810, -0.6681]])
sample_variance(x_dict={})[source]

Return the variance of the distribution.

Parameters:x_dict (dict, defaults to {}) – Parameters of this distribution.

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> var = p1.sample_variance()
>>> print(var)
tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> var = p2.sample_variance({"y": sample_y})
>>> print(var) # doctest: +SKIP
tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])
get_log_prob(x_dict, sum_features=True, feature_dims=None)[source]

Giving variables, this method returns values of log-pdf.

Parameters:
  • x_dict (dict) – Input variables.
  • sum_features (bool, defaults to True) – Whether the output is summed across some dimensions which are specified by feature_dims.
  • feature_dims (list or NoneType, defaults to None) – Set dimensions to sum across the output.
Returns:

log_prob – Values of log-probability density/mass function.

Return type:

torch.Tensor

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> sample_x = torch.randn(1, 10) # Psuedo data
>>> log_prob = p1.log_prob({"x": sample_x})
>>> print(log_prob) # doctest: +SKIP
tensor([-16.1153])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> log_prob = p2.log_prob({"x": sample_x, "y": sample_y})
>>> print(log_prob) # doctest: +SKIP
tensor([-21.5251])
get_entropy(x_dict={}, sum_features=True, feature_dims=None)[source]

Giving variables, this method returns values of entropy.

Parameters:
  • x_dict (dict, defaults to {}) – Input variables.
  • sum_features (bool, defaults to True) – Whether the output is summed across some dimensions which are specified by feature_dims.
  • feature_dims (list or NoneType, defaults to None) – Set dimensions to sum across the output.
Returns:

entropy – Values of entropy.

Return type:

torch.Tensor

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> entropy = p1.get_entropy()
>>> print(entropy)
tensor([14.1894])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> entropy = p2.get_entropy({"y": sample_y})
>>> print(entropy)
tensor([14.1894])
log_prob(sum_features=True, feature_dims=None)[source]

Return an instance of pixyz.losses.LogProb.

Parameters:
  • sum_features (bool, defaults to True) – Whether the output is summed across some axes (dimensions) which are specified by feature_dims.
  • feature_dims (list or NoneType, defaults to None) – Set axes to sum across the output.
Returns:

An instance of pixyz.losses.LogProb

Return type:

pixyz.losses.LogProb

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> sample_x = torch.randn(1, 10) # Psuedo data
>>> log_prob = p1.log_prob().eval({"x": sample_x})
>>> print(log_prob) # doctest: +SKIP
tensor([-16.1153])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> log_prob = p2.log_prob().eval({"x": sample_x, "y": sample_y})
>>> print(log_prob) # doctest: +SKIP
tensor([-21.5251])
prob(sum_features=True, feature_dims=None)[source]

Return an instance of pixyz.losses.LogProb.

Parameters:
  • sum_features (bool, defaults to True) – Choose whether the output is summed across some axes (dimensions) which are specified by feature_dims.
  • feature_dims (list or NoneType, defaults to None) – Set dimensions to sum across the output. (Note: this parameter is not used for now.)
Returns:

An instance of pixyz.losses.Prob

Return type:

pixyz.losses.Prob

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> sample_x = torch.randn(1, 10) # Psuedo data
>>> prob = p1.prob().eval({"x": sample_x})
>>> print(prob) # doctest: +SKIP
tensor([4.0933e-07])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> prob = p2.prob().eval({"x": sample_x, "y": sample_y})
>>> print(prob) # doctest: +SKIP
tensor([2.9628e-09])
forward(*args, **kwargs)[source]

When this class is inherited by DNNs, this method should be overrided.

replace_var(**replace_dict)[source]

Return an instance of pixyz.distributions.ReplaceVarDistribution.

Parameters:replace_dict (dict) – Dictionary.
Returns:An instance of pixyz.distributions.ReplaceVarDistribution
Return type:pixyz.distributions.ReplaceVarDistribution
marginalize_var(marginalize_list)[source]

Return an instance of pixyz.distributions.MarginalizeVarDistribution.

Parameters:marginalize_list (list or other) – Variables to marginalize.
Returns:An instance of pixyz.distributions.MarginalizeVarDistribution
Return type:pixyz.distributions.MarginalizeVarDistribution
extra_repr()[source]

Set the extra representation of the module

To print customized extra information, you should reimplement this method in your own modules. Both single-line and multi-line strings are acceptable.

Exponential families

Normal

class pixyz.distributions.Normal(cond_var=[], var=['x'], name='p', features_shape=torch.Size([]), **kwargs)[source]

Bases: pixyz.distributions.distributions.DistributionBase

Normal distribution parameterized by loc and scale.

params_keys

Return the list of parameter names for this distribution.

Type:list
distribution_torch_class

Return the class of PyTorch distribution.

distribution_name

Name of this distribution class.

Type:str

Laplace

class pixyz.distributions.Laplace(cond_var=[], var=['x'], name='p', features_shape=torch.Size([]), **kwargs)[source]

Bases: pixyz.distributions.distributions.DistributionBase

Laplace distribution parameterized by loc and scale.

params_keys

Return the list of parameter names for this distribution.

Type:list
distribution_torch_class

Return the class of PyTorch distribution.

distribution_name

Name of this distribution class.

Type:str

Bernoulli

class pixyz.distributions.Bernoulli(cond_var=[], var=['x'], name='p', features_shape=torch.Size([]), **kwargs)[source]

Bases: pixyz.distributions.distributions.DistributionBase

Bernoulli distribution parameterized by probs.

params_keys

Return the list of parameter names for this distribution.

Type:list
distribution_torch_class

Return the class of PyTorch distribution.

distribution_name

Name of this distribution class.

Type:str

RelaxedBernoulli

class pixyz.distributions.RelaxedBernoulli(temperature=tensor(0.1000), cond_var=[], var=['x'], name='p', features_shape=torch.Size([]), **kwargs)[source]

Bases: pixyz.distributions.exponential_distributions.Bernoulli

Relaxed (re-parameterizable) Bernoulli distribution parameterized by probs.

temperature
distribution_torch_class

Return the class of PyTorch distribution.

relaxed_distribution_torch_class

Use relaxed version only when sampling

distribution_name

Name of this distribution class.

Type:str
set_dist(x_dict={}, sampling=True, batch_n=None, **kwargs)[source]

Set dist as PyTorch distributions given parameters.

This requires that params_keys and distribution_torch_class are set.

Parameters:
  • x_dict (dict, defaults to {}.) – Parameters of this distribution.
  • sampling (bool, defaults to False.) – Choose whether to use relaxed_* in PyTorch distribution.
  • batch_n (int, defaults to None.) – Set batch size of parameters.
  • **kwargs – Arbitrary keyword arguments.

FactorizedBernoulli

class pixyz.distributions.FactorizedBernoulli(cond_var=[], var=['x'], name='p', features_shape=torch.Size([]), **kwargs)[source]

Bases: pixyz.distributions.exponential_distributions.Bernoulli

Factorized Bernoulli distribution parameterized by probs.

References

[Vedantam+ 2017] Generative Models of Visually Grounded Imagination

distribution_name

Name of this distribution class.

Type:str
get_log_prob(x_dict)[source]

Giving variables, this method returns values of log-pdf.

Parameters:
  • x_dict (dict) – Input variables.
  • sum_features (bool, defaults to True) – Whether the output is summed across some dimensions which are specified by feature_dims.
  • feature_dims (list or NoneType, defaults to None) – Set dimensions to sum across the output.
Returns:

log_prob – Values of log-probability density/mass function.

Return type:

torch.Tensor

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> sample_x = torch.randn(1, 10) # Psuedo data
>>> log_prob = p1.log_prob({"x": sample_x})
>>> print(log_prob) # doctest: +SKIP
tensor([-16.1153])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> log_prob = p2.log_prob({"x": sample_x, "y": sample_y})
>>> print(log_prob) # doctest: +SKIP
tensor([-21.5251])

Categorical

class pixyz.distributions.Categorical(cond_var=[], var=['x'], name='p', features_shape=torch.Size([]), **kwargs)[source]

Bases: pixyz.distributions.distributions.DistributionBase

Categorical distribution parameterized by probs.

params_keys

Return the list of parameter names for this distribution.

Type:list
distribution_torch_class

Return the class of PyTorch distribution.

distribution_name

Name of this distribution class.

Type:str

RelaxedCategorical

class pixyz.distributions.RelaxedCategorical(temperature=tensor(0.1000), cond_var=[], var=['x'], name='p', features_shape=torch.Size([]), **kwargs)[source]

Bases: pixyz.distributions.exponential_distributions.Categorical

Relaxed (re-parameterizable) categorical distribution parameterized by probs.

temperature
distribution_torch_class

Return the class of PyTorch distribution.

relaxed_distribution_torch_class

Use relaxed version only when sampling

distribution_name

Name of this distribution class.

Type:str
set_dist(x_dict={}, sampling=True, batch_n=None, **kwargs)[source]

Set dist as PyTorch distributions given parameters.

This requires that params_keys and distribution_torch_class are set.

Parameters:
  • x_dict (dict, defaults to {}.) – Parameters of this distribution.
  • sampling (bool, defaults to False.) – Choose whether to use relaxed_* in PyTorch distribution.
  • batch_n (int, defaults to None.) – Set batch size of parameters.
  • **kwargs – Arbitrary keyword arguments.
sample_mean(x_dict={})[source]

Return the mean of the distribution.

Parameters:x_dict (dict, defaults to {}) – Parameters of this distribution.

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> mean = p1.sample_mean()
>>> print(mean)
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> mean = p2.sample_mean({"y": sample_y})
>>> print(mean) # doctest: +SKIP
tensor([[-0.2189, -1.0310, -0.1917, -0.3085,  1.5190, -0.9037,  1.2559,  0.1410,
          1.2810, -0.6681]])
sample_variance(x_dict={})[source]

Return the variance of the distribution.

Parameters:x_dict (dict, defaults to {}) – Parameters of this distribution.

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> var = p1.sample_variance()
>>> print(var)
tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> var = p2.sample_variance({"y": sample_y})
>>> print(var) # doctest: +SKIP
tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])

Beta

class pixyz.distributions.Beta(cond_var=[], var=['x'], name='p', features_shape=torch.Size([]), **kwargs)[source]

Bases: pixyz.distributions.distributions.DistributionBase

Beta distribution parameterized by concentration1 and concentration0.

params_keys

Return the list of parameter names for this distribution.

Type:list
distribution_torch_class

Return the class of PyTorch distribution.

distribution_name

Name of this distribution class.

Type:str

Dirichlet

class pixyz.distributions.Dirichlet(cond_var=[], var=['x'], name='p', features_shape=torch.Size([]), **kwargs)[source]

Bases: pixyz.distributions.distributions.DistributionBase

Dirichlet distribution parameterized by concentration.

params_keys

Return the list of parameter names for this distribution.

Type:list
distribution_torch_class

Return the class of PyTorch distribution.

distribution_name

Name of this distribution class.

Type:str

Gamma

class pixyz.distributions.Gamma(cond_var=[], var=['x'], name='p', features_shape=torch.Size([]), **kwargs)[source]

Bases: pixyz.distributions.distributions.DistributionBase

Gamma distribution parameterized by concentration and rate.

params_keys

Return the list of parameter names for this distribution.

Type:list
distribution_torch_class

Return the class of PyTorch distribution.

distribution_name

Name of this distribution class.

Type:str

Complex distributions

MixtureModel

class pixyz.distributions.MixtureModel(distributions, prior, name='p')[source]

Bases: pixyz.distributions.distributions.Distribution

Mixture models.

p(x) = \sum_i p(x|z=i)p(z=i)

Examples

>>> from pixyz.distributions import Normal, Categorical
>>> from pixyz.distributions.mixture_distributions import MixtureModel
>>> z_dim = 3  # the number of mixture
>>> x_dim = 2  # the input dimension.
>>> distributions = []  # the list of distributions
>>> for i in range(z_dim):
...     loc = torch.randn(x_dim)  # initialize the value of location (mean)
...     scale = torch.empty(x_dim).fill_(1.)  # initialize the value of scale (variance)
...     distributions.append(Normal(loc=loc, scale=scale, var=["x"], name="p_%d" %i))
>>> probs = torch.empty(z_dim).fill_(1. / z_dim)  # initialize the value of probabilities
>>> prior = Categorical(probs=probs, var=["z"], name="prior")
>>> p = MixtureModel(distributions=distributions, prior=prior)
>>> print(p)
Distribution:
  p(x) = p_{0}(x|z=0)prior(z=0) + p_{1}(x|z=1)prior(z=1) + p_{2}(x|z=2)prior(z=2)
Network architecture:
  MixtureModel(
    name=p, distribution_name=Mixture Model,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([])
    (distributions): ModuleList(
      (0): Normal(
        name=p_{0}, distribution_name=Normal,
        var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([2])
        (loc): torch.Size([1, 2])
        (scale): torch.Size([1, 2])
      )
      (1): Normal(
        name=p_{1}, distribution_name=Normal,
        var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([2])
        (loc): torch.Size([1, 2])
        (scale): torch.Size([1, 2])
      )
      (2): Normal(
        name=p_{2}, distribution_name=Normal,
        var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([2])
        (loc): torch.Size([1, 2])
        (scale): torch.Size([1, 2])
      )
    )
    (prior): Categorical(
      name=prior, distribution_name=Categorical,
      var=['z'], cond_var=[], input_var=[], features_shape=torch.Size([3])
      (probs): torch.Size([1, 3])
    )
  )
__init__(distributions, prior, name='p')[source]
Parameters:
  • distributions (list) – List of distributions.
  • prior (pixyz.Distribution.Categorical) – Prior distribution of latent variable (i.e., a contribution rate). This should be a categorical distribution and the number of its category should be the same as the length of distributions.
  • name (str, defaults to “p”) – Name of this distribution. This name is displayed in prob_text and prob_factorized_text.
hidden_var

Hidden variables of this distribution.

Type:list
prob_text

Return a formula of the (joint) probability distribution.

Type:str
prob_factorized_text

Return a formula of the factorized probability distribution.

Type:str
distribution_name

Name of this distribution class.

Type:str
posterior(name=None)[source]
sample(batch_n=None, sample_shape=torch.Size([]), return_hidden=False, **kwargs)[source]

Sample variables of this distribution. If cond_var is not empty, you should set inputs as dict.

Parameters:
  • x_dict (torch.Tensor, list, or dict, defaults to {}) – Input variables.
  • sample_shape (list or NoneType, defaults to torch.Size()) – Shape of generating samples.
  • batch_n (int, defaults to None.) – Set batch size of parameters.
  • return_all (bool, defaults to True) – Choose whether the output contains input variables.
  • reparam (bool, defaults to False.) – Choose whether we sample variables with re-parameterized trick.
Returns:

output – Samples of this distribution.

Return type:

dict

Examples

>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...            features_shape=[10, 2])
>>> print(p)
Distribution:
  p(x)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([10, 2])
    (loc): torch.Size([1, 10, 2])
    (scale): torch.Size([1, 10, 2])
  )
>>> p.sample()["x"].shape  # (batch_n=1, features_shape)
torch.Size([1, 10, 2])
>>> p.sample(batch_n=20)["x"].shape  # (batch_n, features_shape)
torch.Size([20, 10, 2])
>>> p.sample(batch_n=20, sample_shape=[40, 30])["x"].shape  # (sample_shape, batch_n, features_shape)
torch.Size([40, 30, 20, 10, 2])
>>> # Conditional distribution
>>> p = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...            features_shape=[10])
>>> print(p)
Distribution:
  p(x|y)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=['y'], input_var=['y'], features_shape=torch.Size([10])
    (scale): torch.Size([1, 10])
  )
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> sample_a = torch.randn(1, 10) # Psuedo data
>>> sample = p.sample({"y": sample_y})
>>> print(sample) # input_var + var  # doctest: +SKIP
{'y': tensor([[-0.5182,  0.3484,  0.9042,  0.1914,  0.6905,
               -1.0859, -0.4433, -0.0255, 0.8198,  0.4571]]),
 'x': tensor([[-0.7205, -1.3996,  0.5528, -0.3059,  0.5384,
               -1.4976, -0.1480,  0.0841,0.3321,  0.5561]])}
>>> sample = p.sample({"y": sample_y, "a": sample_a}) # Redundant input ("a")
>>> print(sample) # input_var + var + "a" (redundant input)  # doctest: +SKIP
{'y': tensor([[ 1.3582, -1.1151, -0.8111,  1.0630,  1.1633,
                0.3855,  2.6324, -0.9357, -0.8649, -0.6015]]),
 'a': tensor([[-0.1874,  1.7958, -1.4084, -2.5646,  1.0868,
               -0.7523, -0.0852, -2.4222, -0.3914, -0.9755]]),
 'x': tensor([[-0.3272, -0.5222, -1.3659,  1.8386,  2.3204,
                0.3686,  0.6311, -1.1208, 0.3656, -0.6683]])}
get_log_prob(x_dict, return_hidden=False, **kwargs)[source]

Evaluate log-pdf, log p(x) (if return_hidden=False) or log p(x, z) (if return_hidden=True).

Parameters:
  • x_dict (dict) – Input variables (including var).
  • return_hidden (bool, defaults to False) –
Returns:

log_prob – The log-pdf value of x.

return_hidden = 0 :

dim=0 : the size of batch

return_hidden = 1 :

dim=0 : the number of mixture

dim=1 : the size of batch

Return type:

torch.Tensor

ProductOfNormal

class pixyz.distributions.ProductOfNormal(p=[], name='p', features_shape=torch.Size([]))[source]

Bases: pixyz.distributions.exponential_distributions.Normal

Product of normal distributions.

p(z|x,y) \propto p(z)p(z|x)p(z|y)

In this model, p(z|x) and p(a|y) perform as experts and p(z) corresponds a prior of experts.

References

[Vedantam+ 2017] Generative Models of Visually Grounded Imagination

[Wu+ 2018] Multimodal Generative Models for Scalable Weakly-Supervised Learning

Examples

>>> pon = ProductOfNormal([p_x, p_y]) # doctest: +SKIP
>>> pon.sample({"x": x, "y": y}) # doctest: +SKIP
{'x': tensor([[0., 0., 0.,  ..., 0., 0., 0.],
     [0., 0., 0.,  ..., 0., 0., 0.],
     [0., 0., 0.,  ..., 0., 0., 0.],
     ...,
     [0., 0., 0.,  ..., 0., 0., 0.],
     [0., 0., 0.,  ..., 0., 0., 0.],
     [0., 0., 0.,  ..., 0., 0., 0.]],),
 'y': tensor([[0., 0., 0.,  ..., 0., 0., 1.],
     [0., 0., 1.,  ..., 0., 0., 0.],
     [0., 1., 0.,  ..., 0., 0., 0.],
     ...,
     [0., 0., 0.,  ..., 0., 1., 0.],
     [1., 0., 0.,  ..., 0., 0., 0.],
     [0., 0., 0.,  ..., 0., 0., 1.]]),
 'z': tensor([[ 0.6611,  0.3811,  0.7778,  ..., -0.0468, -0.3615, -0.6569],
     [-0.0071, -0.9178,  0.6620,  ..., -0.1472,  0.6023,  0.5903],
     [-0.3723, -0.7758,  0.0195,  ...,  0.8239, -0.3537,  0.3854],
     ...,
     [ 0.7820, -0.4761,  0.1804,  ..., -0.5701, -0.0714, -0.5485],
     [-0.1873, -0.2105, -0.1861,  ..., -0.5372,  0.0752,  0.2777],
     [-0.2563, -0.0828,  0.1605,  ...,  0.2767, -0.8456,  0.7364]])}
>>> pon.sample({"y": y}) # doctest: +SKIP
{'y': tensor([[0., 0., 0.,  ..., 0., 0., 0.],
     [0., 0., 0.,  ..., 0., 0., 1.],
     [0., 0., 0.,  ..., 1., 0., 0.],
     ...,
     [0., 0., 0.,  ..., 0., 0., 0.],
     [0., 1., 0.,  ..., 0., 0., 0.],
     [0., 0., 0.,  ..., 0., 0., 0.]]),
 'z': tensor([[-0.3264, -0.4448,  0.3610,  ..., -0.7378,  0.3002,  0.4370],
     [ 0.0928, -0.1830,  1.1768,  ...,  1.1808, -0.7226, -0.4152],
     [ 0.6999,  0.2222, -0.2901,  ...,  0.5706,  0.7091,  0.5179],
     ...,
     [ 0.5688, -1.6612, -0.0713,  ..., -0.1400, -0.3903,  0.2533],
     [ 0.5412, -0.0289,  0.6365,  ...,  0.7407,  0.7838,  0.9218],
     [ 0.0299,  0.5148, -0.1001,  ...,  0.9938,  1.0689, -1.1902]])}
>>> pon.sample()  # same as sampling from unit Gaussian. # doctest: +SKIP
{'z': tensor(-0.4494)}
__init__(p=[], name='p', features_shape=torch.Size([]))[source]
Parameters:
  • p (list of pixyz.distributions.Normal.) – List of experts.
  • name (str, defaults to “p”) – Name of this distribution. This name is displayed in prob_text and prob_factorized_text.
  • features_shape (torch.Size or list, defaults to torch.Size())) – Shape of dimensions (features) of this distribution.
prob_factorized_text

Return a formula of the factorized probability distribution.

Type:str
prob_joint_factorized_and_text

Return a formula of the factorized probability distribution.

Type:str
get_params(params_dict={}, **kwargs)[source]

This method aims to get parameters of this distributions from constant parameters set in initialization and outputs of DNNs.

Parameters:params_dict (dict, defaults to {}) – Input parameters.
Returns:output_dict – Output parameters.
Return type:dict

Examples

>>> from pixyz.distributions import Normal
>>> dist_1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...                 features_shape=[1])
>>> print(dist_1)
Distribution:
  p(x)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([1])
    (loc): torch.Size([1, 1])
    (scale): torch.Size([1, 1])
  )
>>> dist_1.get_params()
{'loc': tensor([[0.]]), 'scale': tensor([[1.]])}
>>> dist_2 = Normal(loc=torch.tensor(0.), scale="z", cond_var=["z"], var=["x"])
>>> print(dist_2)
Distribution:
  p(x|z)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
    (loc): torch.Size([1])
  )
>>> dist_2.get_params({"z": torch.tensor(1.)})
{'scale': tensor(1.), 'loc': tensor([0.])}
log_prob(sum_features=True, feature_dims=None)[source]

Return an instance of pixyz.losses.LogProb.

Parameters:
  • sum_features (bool, defaults to True) – Whether the output is summed across some axes (dimensions) which are specified by feature_dims.
  • feature_dims (list or NoneType, defaults to None) – Set axes to sum across the output.
Returns:

An instance of pixyz.losses.LogProb

Return type:

pixyz.losses.LogProb

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> sample_x = torch.randn(1, 10) # Psuedo data
>>> log_prob = p1.log_prob().eval({"x": sample_x})
>>> print(log_prob) # doctest: +SKIP
tensor([-16.1153])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> log_prob = p2.log_prob().eval({"x": sample_x, "y": sample_y})
>>> print(log_prob) # doctest: +SKIP
tensor([-21.5251])
prob(sum_features=True, feature_dims=None)[source]

Return an instance of pixyz.losses.LogProb.

Parameters:
  • sum_features (bool, defaults to True) – Choose whether the output is summed across some axes (dimensions) which are specified by feature_dims.
  • feature_dims (list or NoneType, defaults to None) – Set dimensions to sum across the output. (Note: this parameter is not used for now.)
Returns:

An instance of pixyz.losses.Prob

Return type:

pixyz.losses.Prob

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> sample_x = torch.randn(1, 10) # Psuedo data
>>> prob = p1.prob().eval({"x": sample_x})
>>> print(prob) # doctest: +SKIP
tensor([4.0933e-07])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> prob = p2.prob().eval({"x": sample_x, "y": sample_y})
>>> print(prob) # doctest: +SKIP
tensor([2.9628e-09])
get_log_prob(x_dict, sum_features=True, feature_dims=None)[source]

Giving variables, this method returns values of log-pdf.

Parameters:
  • x_dict (dict) – Input variables.
  • sum_features (bool, defaults to True) – Whether the output is summed across some dimensions which are specified by feature_dims.
  • feature_dims (list or NoneType, defaults to None) – Set dimensions to sum across the output.
Returns:

log_prob – Values of log-probability density/mass function.

Return type:

torch.Tensor

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> sample_x = torch.randn(1, 10) # Psuedo data
>>> log_prob = p1.log_prob({"x": sample_x})
>>> print(log_prob) # doctest: +SKIP
tensor([-16.1153])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> log_prob = p2.log_prob({"x": sample_x, "y": sample_y})
>>> print(log_prob) # doctest: +SKIP
tensor([-21.5251])

ElementWiseProductOfNormal

class pixyz.distributions.ElementWiseProductOfNormal(p, name='p', features_shape=torch.Size([]))[source]

Bases: pixyz.distributions.poe.ProductOfNormal

Product of normal distributions. In this distribution, each element of the input vector on the given distribution is considered as a different expert.

p(z|x) = p(z|x_1, x_2) \propto p(z)p(z|x_1)p(z|x_2)

Examples

>>> pon = ElementWiseProductOfNormal(p) # doctest: +SKIP
>>> pon.sample({"x": x}) # doctest: +SKIP
{'x': tensor([[0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
     [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.]]),
 'z': tensor([[-0.3572, -0.0632,  0.4872,  0.2269, -0.1693, -0.0160, -0.0429,  0.2017,
      -0.1589, -0.3380, -0.9598,  0.6216, -0.4296, -1.1349,  0.0901,  0.3994,
       0.2313, -0.5227, -0.7973,  0.3968,  0.7137, -0.5639, -0.4891, -0.1249,
       0.8256,  0.1463,  0.0801, -1.2202,  0.6984, -0.4036,  0.4960, -0.4376,
       0.3310, -0.2243, -0.2381, -0.2200,  0.8969,  0.2674,  0.4681,  1.6764,
       0.8127,  0.2722, -0.2048,  0.1903, -0.1398,  0.0099,  0.4382, -0.8016,
       0.9947,  0.7556, -0.2017, -0.3920,  1.4212, -1.2529, -0.1002, -0.0031,
       0.1876,  0.4267,  0.3622,  0.2648,  0.4752,  0.0843, -0.3065, -0.4922],
     [ 0.3770, -0.0413,  0.9102,  0.2897, -0.0567,  0.5211,  1.5233, -0.3539,
       0.5163, -0.2271, -0.1027,  0.0294, -1.4617,  0.1640,  0.2025, -0.2190,
       0.0555,  0.5779, -0.2930, -0.2161,  0.2835, -0.0354, -0.2569, -0.7171,
       0.0164, -0.4080,  1.1088,  0.3947,  0.2720, -0.0600, -0.9295, -0.0234,
       0.5624,  0.4866,  0.5285,  1.1827,  0.2494,  0.0777,  0.7585,  0.5127,
       0.7500, -0.3253,  0.0250,  0.0888,  1.0340, -0.1405, -0.8114,  0.4492,
       0.2725, -0.0270,  0.6379, -0.8096,  0.4259,  0.3179, -0.1681,  0.3365,
       0.6305,  0.5203,  0.2384,  0.0572,  0.4804,  0.9553, -0.3244,  1.5373]])}
>>> pon.sample({"x": torch.zeros_like(x)})  # same as sampling from unit Gaussian. # doctest: +SKIP
{'x': tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
     [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]),
 'z': tensor([[-0.7777, -0.5908, -1.5498, -0.7505,  0.6201,  0.7218,  1.0045,  0.8923,
      -0.8030, -0.3569,  0.2932,  0.2122,  0.1640,  0.7893, -0.3500, -1.0537,
      -1.2769,  0.6122, -1.0083, -0.2915, -0.1928, -0.7486,  0.2418, -1.9013,
       1.2514,  1.3035, -0.3029, -0.3098, -0.5415,  1.1970, -0.4443,  2.2393,
      -0.6980,  0.2820,  1.6972,  0.6322,  0.4308,  0.8953,  0.7248,  0.4440,
       2.2770,  1.7791,  0.7563, -1.1781, -0.8331,  0.1825,  1.5447,  0.1385,
      -1.1348,  0.0257,  0.3374,  0.5889,  1.1231, -1.2476, -0.3801, -1.4404,
      -1.3066, -1.2653,  0.5958, -1.7423,  0.7189, -0.7236,  0.2330,  0.3117],
     [ 0.5495,  0.7210, -0.4708, -2.0631, -0.6170,  0.2436, -0.0133, -0.4616,
      -0.8091, -0.1592,  1.3117,  0.0276,  0.6625, -0.3748, -0.5049,  1.8260,
      -0.3631,  1.1546, -1.0913,  0.2712,  1.5493,  1.4294, -2.1245, -2.0422,
       0.4976, -1.2785,  0.5028,  1.4240,  1.1983,  0.2468,  1.1682, -0.6725,
      -1.1198, -1.4942, -0.3629,  0.1325, -0.2256,  0.4280,  0.9830, -1.9427,
      -0.2181,  1.1850, -0.7514, -0.8172,  2.1031, -0.1698, -0.3777, -0.7863,
       1.0936, -1.3720,  0.9999,  1.3302, -0.8954, -0.5999,  2.3305,  0.5702,
      -1.0767, -0.2750, -0.3741, -0.7026, -1.5408,  0.0667,  1.2550, -0.5117]])}
__init__(p, name='p', features_shape=torch.Size([]))[source]
Parameters:
  • p (pixyz.distributions.Normal) – Each element of this input vector is considered as a different expert. When some elements are 0, experts corresponding to these elements are considered not to be specified. p(z|x) = p(z|x_1, x_2=0) \propto p(z)p(z|x_1)
  • name (str, defaults to "p") – Name of this distribution. This name is displayed in prob_text and prob_factorized_text.
  • features_shape (torch.Size or list, defaults to torch.Size())) – Shape of dimensions (features) of this distribution.

Flow distributions

TransformedDistribution

class pixyz.distributions.TransformedDistribution(prior, flow, var, name='p')[source]

Bases: pixyz.distributions.distributions.Distribution

Convert flow transformations to distributions.

p(z=f_{flow}(x)),

where x \sim p_{prior}(x).

Once initializing, it can be handled as a distribution module.

distribution_name

Name of this distribution class.

Type:str
flow_input_var

Input variables of the flow module.

Type:list
prob_factorized_text

Return a formula of the factorized probability distribution.

Type:str
logdet_jacobian

Get log-determinant Jacobian.

Before calling this, you should run forward or update_jacobian methods to calculate and store log-determinant Jacobian.

sample(x_dict={}, batch_n=None, sample_shape=torch.Size([]), return_all=True, reparam=False, compute_jacobian=True)[source]

Sample variables of this distribution. If cond_var is not empty, you should set inputs as dict.

Parameters:
  • x_dict (torch.Tensor, list, or dict, defaults to {}) – Input variables.
  • sample_shape (list or NoneType, defaults to torch.Size()) – Shape of generating samples.
  • batch_n (int, defaults to None.) – Set batch size of parameters.
  • return_all (bool, defaults to True) – Choose whether the output contains input variables.
  • reparam (bool, defaults to False.) – Choose whether we sample variables with re-parameterized trick.
Returns:

output – Samples of this distribution.

Return type:

dict

Examples

>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...            features_shape=[10, 2])
>>> print(p)
Distribution:
  p(x)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([10, 2])
    (loc): torch.Size([1, 10, 2])
    (scale): torch.Size([1, 10, 2])
  )
>>> p.sample()["x"].shape  # (batch_n=1, features_shape)
torch.Size([1, 10, 2])
>>> p.sample(batch_n=20)["x"].shape  # (batch_n, features_shape)
torch.Size([20, 10, 2])
>>> p.sample(batch_n=20, sample_shape=[40, 30])["x"].shape  # (sample_shape, batch_n, features_shape)
torch.Size([40, 30, 20, 10, 2])
>>> # Conditional distribution
>>> p = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...            features_shape=[10])
>>> print(p)
Distribution:
  p(x|y)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=['y'], input_var=['y'], features_shape=torch.Size([10])
    (scale): torch.Size([1, 10])
  )
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> sample_a = torch.randn(1, 10) # Psuedo data
>>> sample = p.sample({"y": sample_y})
>>> print(sample) # input_var + var  # doctest: +SKIP
{'y': tensor([[-0.5182,  0.3484,  0.9042,  0.1914,  0.6905,
               -1.0859, -0.4433, -0.0255, 0.8198,  0.4571]]),
 'x': tensor([[-0.7205, -1.3996,  0.5528, -0.3059,  0.5384,
               -1.4976, -0.1480,  0.0841,0.3321,  0.5561]])}
>>> sample = p.sample({"y": sample_y, "a": sample_a}) # Redundant input ("a")
>>> print(sample) # input_var + var + "a" (redundant input)  # doctest: +SKIP
{'y': tensor([[ 1.3582, -1.1151, -0.8111,  1.0630,  1.1633,
                0.3855,  2.6324, -0.9357, -0.8649, -0.6015]]),
 'a': tensor([[-0.1874,  1.7958, -1.4084, -2.5646,  1.0868,
               -0.7523, -0.0852, -2.4222, -0.3914, -0.9755]]),
 'x': tensor([[-0.3272, -0.5222, -1.3659,  1.8386,  2.3204,
                0.3686,  0.6311, -1.1208, 0.3656, -0.6683]])}
get_log_prob(x_dict, sum_features=True, feature_dims=None, compute_jacobian=False)[source]

Giving variables, this method returns values of log-pdf.

Parameters:
  • x_dict (dict) – Input variables.
  • sum_features (bool, defaults to True) – Whether the output is summed across some dimensions which are specified by feature_dims.
  • feature_dims (list or NoneType, defaults to None) – Set dimensions to sum across the output.
Returns:

log_prob – Values of log-probability density/mass function.

Return type:

torch.Tensor

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> sample_x = torch.randn(1, 10) # Psuedo data
>>> log_prob = p1.log_prob({"x": sample_x})
>>> print(log_prob) # doctest: +SKIP
tensor([-16.1153])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> log_prob = p2.log_prob({"x": sample_x, "y": sample_y})
>>> print(log_prob) # doctest: +SKIP
tensor([-21.5251])
forward(x, y=None, compute_jacobian=True)[source]

Forward propagation of flow layers.

Parameters:
  • x (torch.Tensor) – Input data.
  • y (torch.Tensor, defaults to None) – Data for conditioning.
  • compute_jacobian (bool, defaults to True) – Whether to calculate and store log-determinant Jacobian. If true, calculated Jacobian values are stored in logdet_jacobian.
Returns:

z

Return type:

torch.Tensor

inverse(z, y=None)[source]

Backward (inverse) propagation of flow layers. In this method, log-determinant Jacobian is not calculated.

Parameters:
  • z (torch.Tensor) – Input data.
  • y (torch.Tensor, defaults to None) – Data for conditioning.
Returns:

x

Return type:

torch.Tensor

InverseTransformedDistribution

class pixyz.distributions.InverseTransformedDistribution(prior, flow, var, cond_var=[], name='p')[source]

Bases: pixyz.distributions.distributions.Distribution

Convert inverse flow transformations to distributions.

p(x=f^{-1}_{flow}(z)),

where z \sim p_{prior}(z).

Once initializing, it can be handled as a distribution module.

Moreover, this distribution can take a conditional variable.

p(x=f^{-1}_{flow}(z, y)),

where z \sim p_{prior}(z) and y is given.

distribution_name

Name of this distribution class.

Type:str
flow_output_var
prob_factorized_text

Return a formula of the factorized probability distribution.

Type:str
logdet_jacobian

Get log-determinant Jacobian.

Before calling this, you should run forward or update_jacobian methods to calculate and store log-determinant Jacobian.

sample(x_dict={}, batch_n=None, sample_shape=torch.Size([]), return_all=True, reparam=False)[source]

Sample variables of this distribution. If cond_var is not empty, you should set inputs as dict.

Parameters:
  • x_dict (torch.Tensor, list, or dict, defaults to {}) – Input variables.
  • sample_shape (list or NoneType, defaults to torch.Size()) – Shape of generating samples.
  • batch_n (int, defaults to None.) – Set batch size of parameters.
  • return_all (bool, defaults to True) – Choose whether the output contains input variables.
  • reparam (bool, defaults to False.) – Choose whether we sample variables with re-parameterized trick.
Returns:

output – Samples of this distribution.

Return type:

dict

Examples

>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...            features_shape=[10, 2])
>>> print(p)
Distribution:
  p(x)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([10, 2])
    (loc): torch.Size([1, 10, 2])
    (scale): torch.Size([1, 10, 2])
  )
>>> p.sample()["x"].shape  # (batch_n=1, features_shape)
torch.Size([1, 10, 2])
>>> p.sample(batch_n=20)["x"].shape  # (batch_n, features_shape)
torch.Size([20, 10, 2])
>>> p.sample(batch_n=20, sample_shape=[40, 30])["x"].shape  # (sample_shape, batch_n, features_shape)
torch.Size([40, 30, 20, 10, 2])
>>> # Conditional distribution
>>> p = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...            features_shape=[10])
>>> print(p)
Distribution:
  p(x|y)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=['y'], input_var=['y'], features_shape=torch.Size([10])
    (scale): torch.Size([1, 10])
  )
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> sample_a = torch.randn(1, 10) # Psuedo data
>>> sample = p.sample({"y": sample_y})
>>> print(sample) # input_var + var  # doctest: +SKIP
{'y': tensor([[-0.5182,  0.3484,  0.9042,  0.1914,  0.6905,
               -1.0859, -0.4433, -0.0255, 0.8198,  0.4571]]),
 'x': tensor([[-0.7205, -1.3996,  0.5528, -0.3059,  0.5384,
               -1.4976, -0.1480,  0.0841,0.3321,  0.5561]])}
>>> sample = p.sample({"y": sample_y, "a": sample_a}) # Redundant input ("a")
>>> print(sample) # input_var + var + "a" (redundant input)  # doctest: +SKIP
{'y': tensor([[ 1.3582, -1.1151, -0.8111,  1.0630,  1.1633,
                0.3855,  2.6324, -0.9357, -0.8649, -0.6015]]),
 'a': tensor([[-0.1874,  1.7958, -1.4084, -2.5646,  1.0868,
               -0.7523, -0.0852, -2.4222, -0.3914, -0.9755]]),
 'x': tensor([[-0.3272, -0.5222, -1.3659,  1.8386,  2.3204,
                0.3686,  0.6311, -1.1208, 0.3656, -0.6683]])}
inference(x_dict, return_all=True, compute_jacobian=False)[source]
get_log_prob(x_dict, sum_features=True, feature_dims=None)[source]

Giving variables, this method returns values of log-pdf.

Parameters:
  • x_dict (dict) – Input variables.
  • sum_features (bool, defaults to True) – Whether the output is summed across some dimensions which are specified by feature_dims.
  • feature_dims (list or NoneType, defaults to None) – Set dimensions to sum across the output.
Returns:

log_prob – Values of log-probability density/mass function.

Return type:

torch.Tensor

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> sample_x = torch.randn(1, 10) # Psuedo data
>>> log_prob = p1.log_prob({"x": sample_x})
>>> print(log_prob) # doctest: +SKIP
tensor([-16.1153])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> log_prob = p2.log_prob({"x": sample_x, "y": sample_y})
>>> print(log_prob) # doctest: +SKIP
tensor([-21.5251])
forward(x, y=None, compute_jacobian=True)[source]

Forward propagation of flow layers.

Parameters:
  • x (torch.Tensor) – Input data.
  • y (torch.Tensor, defaults to None) – Data for conditioning.
  • compute_jacobian (bool, defaults to True) – Whether to calculate and store log-determinant Jacobian. If true, calculated Jacobian values are stored in logdet_jacobian.
Returns:

z

Return type:

torch.Tensor

inverse(z, y=None)[source]

Backward (inverse) propagation of flow layers. In this method, log-determinant Jacobian is not calculated.

Parameters:
  • z (torch.Tensor) – Input data.
  • y (torch.Tensor, defaults to None) – Data for conditioning.
Returns:

x

Return type:

torch.Tensor

Special distributions

Deterministic

class pixyz.distributions.Deterministic(**kwargs)[source]

Bases: pixyz.distributions.distributions.Distribution

Deterministic distribution (or degeneration distribution)

Examples

>>> import torch
>>> class Generator(Deterministic):
...     def __init__(self):
...         super().__init__(cond_var=["z"], var=["x"])
...         self.model = torch.nn.Linear(64, 512)
...     def forward(self, z):
...         return {"x": self.model(z)}
>>> p = Generator()
>>> print(p)
Distribution:
  p(x|z)
Network architecture:
  Generator(
    name=p, distribution_name=Deterministic,
    var=['x'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
    (model): Linear(in_features=64, out_features=512, bias=True)
  )
>>> sample = p.sample({"z": torch.randn(1, 64)})
>>> p.log_prob().eval(sample) # log_prob is not defined.
Traceback (most recent call last):
 ...
NotImplementedError
distribution_name

Name of this distribution class.

Type:str
sample(x_dict={}, return_all=True, **kwargs)[source]

Sample variables of this distribution. If cond_var is not empty, you should set inputs as dict.

Parameters:
  • x_dict (torch.Tensor, list, or dict, defaults to {}) – Input variables.
  • sample_shape (list or NoneType, defaults to torch.Size()) – Shape of generating samples.
  • batch_n (int, defaults to None.) – Set batch size of parameters.
  • return_all (bool, defaults to True) – Choose whether the output contains input variables.
  • reparam (bool, defaults to False.) – Choose whether we sample variables with re-parameterized trick.
Returns:

output – Samples of this distribution.

Return type:

dict

Examples

>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...            features_shape=[10, 2])
>>> print(p)
Distribution:
  p(x)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([10, 2])
    (loc): torch.Size([1, 10, 2])
    (scale): torch.Size([1, 10, 2])
  )
>>> p.sample()["x"].shape  # (batch_n=1, features_shape)
torch.Size([1, 10, 2])
>>> p.sample(batch_n=20)["x"].shape  # (batch_n, features_shape)
torch.Size([20, 10, 2])
>>> p.sample(batch_n=20, sample_shape=[40, 30])["x"].shape  # (sample_shape, batch_n, features_shape)
torch.Size([40, 30, 20, 10, 2])
>>> # Conditional distribution
>>> p = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...            features_shape=[10])
>>> print(p)
Distribution:
  p(x|y)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=['y'], input_var=['y'], features_shape=torch.Size([10])
    (scale): torch.Size([1, 10])
  )
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> sample_a = torch.randn(1, 10) # Psuedo data
>>> sample = p.sample({"y": sample_y})
>>> print(sample) # input_var + var  # doctest: +SKIP
{'y': tensor([[-0.5182,  0.3484,  0.9042,  0.1914,  0.6905,
               -1.0859, -0.4433, -0.0255, 0.8198,  0.4571]]),
 'x': tensor([[-0.7205, -1.3996,  0.5528, -0.3059,  0.5384,
               -1.4976, -0.1480,  0.0841,0.3321,  0.5561]])}
>>> sample = p.sample({"y": sample_y, "a": sample_a}) # Redundant input ("a")
>>> print(sample) # input_var + var + "a" (redundant input)  # doctest: +SKIP
{'y': tensor([[ 1.3582, -1.1151, -0.8111,  1.0630,  1.1633,
                0.3855,  2.6324, -0.9357, -0.8649, -0.6015]]),
 'a': tensor([[-0.1874,  1.7958, -1.4084, -2.5646,  1.0868,
               -0.7523, -0.0852, -2.4222, -0.3914, -0.9755]]),
 'x': tensor([[-0.3272, -0.5222, -1.3659,  1.8386,  2.3204,
                0.3686,  0.6311, -1.1208, 0.3656, -0.6683]])}
sample_mean(x_dict)[source]

Return the mean of the distribution.

Parameters:x_dict (dict, defaults to {}) – Parameters of this distribution.

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> mean = p1.sample_mean()
>>> print(mean)
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> mean = p2.sample_mean({"y": sample_y})
>>> print(mean) # doctest: +SKIP
tensor([[-0.2189, -1.0310, -0.1917, -0.3085,  1.5190, -0.9037,  1.2559,  0.1410,
          1.2810, -0.6681]])

DataDistribution

class pixyz.distributions.DataDistribution(var, name='p_{data}')[source]

Bases: pixyz.distributions.distributions.Distribution

Data distribution.

Samples from this distribution equal given inputs.

Examples

>>> import torch
>>> p = DataDistribution(var=["x"])
>>> print(p)
Distribution:
  p_{data}(x)
Network architecture:
  DataDistribution(
    name=p_{data}, distribution_name=Data distribution,
    var=['x'], cond_var=[], input_var=['x'], features_shape=torch.Size([])
  )
>>> sample = p.sample({"x": torch.randn(1, 64)})
distribution_name

Name of this distribution class.

Type:str
sample(x_dict={}, **kwargs)[source]

Sample variables of this distribution. If cond_var is not empty, you should set inputs as dict.

Parameters:
  • x_dict (torch.Tensor, list, or dict, defaults to {}) – Input variables.
  • sample_shape (list or NoneType, defaults to torch.Size()) – Shape of generating samples.
  • batch_n (int, defaults to None.) – Set batch size of parameters.
  • return_all (bool, defaults to True) – Choose whether the output contains input variables.
  • reparam (bool, defaults to False.) – Choose whether we sample variables with re-parameterized trick.
Returns:

output – Samples of this distribution.

Return type:

dict

Examples

>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...            features_shape=[10, 2])
>>> print(p)
Distribution:
  p(x)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([10, 2])
    (loc): torch.Size([1, 10, 2])
    (scale): torch.Size([1, 10, 2])
  )
>>> p.sample()["x"].shape  # (batch_n=1, features_shape)
torch.Size([1, 10, 2])
>>> p.sample(batch_n=20)["x"].shape  # (batch_n, features_shape)
torch.Size([20, 10, 2])
>>> p.sample(batch_n=20, sample_shape=[40, 30])["x"].shape  # (sample_shape, batch_n, features_shape)
torch.Size([40, 30, 20, 10, 2])
>>> # Conditional distribution
>>> p = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...            features_shape=[10])
>>> print(p)
Distribution:
  p(x|y)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=['y'], input_var=['y'], features_shape=torch.Size([10])
    (scale): torch.Size([1, 10])
  )
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> sample_a = torch.randn(1, 10) # Psuedo data
>>> sample = p.sample({"y": sample_y})
>>> print(sample) # input_var + var  # doctest: +SKIP
{'y': tensor([[-0.5182,  0.3484,  0.9042,  0.1914,  0.6905,
               -1.0859, -0.4433, -0.0255, 0.8198,  0.4571]]),
 'x': tensor([[-0.7205, -1.3996,  0.5528, -0.3059,  0.5384,
               -1.4976, -0.1480,  0.0841,0.3321,  0.5561]])}
>>> sample = p.sample({"y": sample_y, "a": sample_a}) # Redundant input ("a")
>>> print(sample) # input_var + var + "a" (redundant input)  # doctest: +SKIP
{'y': tensor([[ 1.3582, -1.1151, -0.8111,  1.0630,  1.1633,
                0.3855,  2.6324, -0.9357, -0.8649, -0.6015]]),
 'a': tensor([[-0.1874,  1.7958, -1.4084, -2.5646,  1.0868,
               -0.7523, -0.0852, -2.4222, -0.3914, -0.9755]]),
 'x': tensor([[-0.3272, -0.5222, -1.3659,  1.8386,  2.3204,
                0.3686,  0.6311, -1.1208, 0.3656, -0.6683]])}
sample_mean(x_dict)[source]

Return the mean of the distribution.

Parameters:x_dict (dict, defaults to {}) – Parameters of this distribution.

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> mean = p1.sample_mean()
>>> print(mean)
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> mean = p2.sample_mean({"y": sample_y})
>>> print(mean) # doctest: +SKIP
tensor([[-0.2189, -1.0310, -0.1917, -0.3085,  1.5190, -0.9037,  1.2559,  0.1410,
          1.2810, -0.6681]])
input_var

In DataDistribution, input_var is same as var.

CustomProb

class pixyz.distributions.CustomProb(log_prob_function, var, distribution_name='Custom PDF', **kwargs)[source]

Bases: pixyz.distributions.distributions.Distribution

This distribution is constructed by user-defined probability density/mass function.

Note that this distribution cannot perform sampling.

Examples

>>> import torch
>>> # banana shaped distribution
>>> def log_prob(z):
...     z1, z2 = torch.chunk(z, chunks=2, dim=1)
...     norm = torch.sqrt(z1 ** 2 + z2 ** 2)
...     exp1 = torch.exp(-0.5 * ((z1 - 2) / 0.6) ** 2)
...     exp2 = torch.exp(-0.5 * ((z1 + 2) / 0.6) ** 2)
...     u = 0.5 * ((norm - 2) / 0.4) ** 2 - torch.log(exp1 + exp2)
...     return -u
...
>>> p = CustomProb(log_prob, var=["z"])
>>> loss = p.log_prob().eval({"z": torch.randn(10, 2)})
__init__(log_prob_function, var, distribution_name='Custom PDF', **kwargs)[source]
Parameters:
  • log_prob_function (function) – User-defined log-probability density/mass function.
  • var (list) – Variables of this distribution.
  • distribution_name (str, optional) – Name of this distribution.
  • +*kwargs – Arbitrary keyword arguments.
log_prob_function

User-defined log-probability density/mass function.

input_var

Input variables of this distribution. Normally, it has same values as cond_var.

Type:list
distribution_name

Name of this distribution class.

Type:str
get_log_prob(x_dict, sum_features=True, feature_dims=None)[source]

Giving variables, this method returns values of log-pdf.

Parameters:
  • x_dict (dict) – Input variables.
  • sum_features (bool, defaults to True) – Whether the output is summed across some dimensions which are specified by feature_dims.
  • feature_dims (list or NoneType, defaults to None) – Set dimensions to sum across the output.
Returns:

log_prob – Values of log-probability density/mass function.

Return type:

torch.Tensor

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> sample_x = torch.randn(1, 10) # Psuedo data
>>> log_prob = p1.log_prob({"x": sample_x})
>>> print(log_prob) # doctest: +SKIP
tensor([-16.1153])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> log_prob = p2.log_prob({"x": sample_x, "y": sample_y})
>>> print(log_prob) # doctest: +SKIP
tensor([-21.5251])

Operators

ReplaceVarDistribution

class pixyz.distributions.ReplaceVarDistribution(p, replace_dict)[source]

Bases: pixyz.distributions.distributions.Distribution

Replace names of variables in Distribution.

Examples

>>> p = DistributionBase(var=["x"], cond_var=["z"])
>>> print(p)
Distribution:
  p(x|z)
Network architecture:
  DistributionBase(
    name=p, distribution_name=,
    var=['x'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
  )
>>> replace_dict = {'x': 'y'}
>>> p_repl = ReplaceVarDistribution(p, replace_dict)
>>> print(p_repl)
Distribution:
  p(y|z)
Network architecture:
  ReplaceVarDistribution(
    name=p, distribution_name=,
    var=['y'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
    (p): DistributionBase(
      name=p, distribution_name=,
      var=['x'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
    )
  )
__init__(p, replace_dict)[source]
Parameters:
forward(*args, **kwargs)[source]

When this class is inherited by DNNs, this method should be overrided.

get_params(params_dict={})[source]

This method aims to get parameters of this distributions from constant parameters set in initialization and outputs of DNNs.

Parameters:params_dict (dict, defaults to {}) – Input parameters.
Returns:output_dict – Output parameters.
Return type:dict

Examples

>>> from pixyz.distributions import Normal
>>> dist_1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...                 features_shape=[1])
>>> print(dist_1)
Distribution:
  p(x)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([1])
    (loc): torch.Size([1, 1])
    (scale): torch.Size([1, 1])
  )
>>> dist_1.get_params()
{'loc': tensor([[0.]]), 'scale': tensor([[1.]])}
>>> dist_2 = Normal(loc=torch.tensor(0.), scale="z", cond_var=["z"], var=["x"])
>>> print(dist_2)
Distribution:
  p(x|z)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
    (loc): torch.Size([1])
  )
>>> dist_2.get_params({"z": torch.tensor(1.)})
{'scale': tensor(1.), 'loc': tensor([0.])}
set_dist(x_dict={}, sampling=False, batch_n=None, **kwargs)[source]
sample(x_dict={}, batch_n=None, sample_shape=torch.Size([]), return_all=True, reparam=False)[source]

Sample variables of this distribution. If cond_var is not empty, you should set inputs as dict.

Parameters:
  • x_dict (torch.Tensor, list, or dict, defaults to {}) – Input variables.
  • sample_shape (list or NoneType, defaults to torch.Size()) – Shape of generating samples.
  • batch_n (int, defaults to None.) – Set batch size of parameters.
  • return_all (bool, defaults to True) – Choose whether the output contains input variables.
  • reparam (bool, defaults to False.) – Choose whether we sample variables with re-parameterized trick.
Returns:

output – Samples of this distribution.

Return type:

dict

Examples

>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...            features_shape=[10, 2])
>>> print(p)
Distribution:
  p(x)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([10, 2])
    (loc): torch.Size([1, 10, 2])
    (scale): torch.Size([1, 10, 2])
  )
>>> p.sample()["x"].shape  # (batch_n=1, features_shape)
torch.Size([1, 10, 2])
>>> p.sample(batch_n=20)["x"].shape  # (batch_n, features_shape)
torch.Size([20, 10, 2])
>>> p.sample(batch_n=20, sample_shape=[40, 30])["x"].shape  # (sample_shape, batch_n, features_shape)
torch.Size([40, 30, 20, 10, 2])
>>> # Conditional distribution
>>> p = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...            features_shape=[10])
>>> print(p)
Distribution:
  p(x|y)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=['y'], input_var=['y'], features_shape=torch.Size([10])
    (scale): torch.Size([1, 10])
  )
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> sample_a = torch.randn(1, 10) # Psuedo data
>>> sample = p.sample({"y": sample_y})
>>> print(sample) # input_var + var  # doctest: +SKIP
{'y': tensor([[-0.5182,  0.3484,  0.9042,  0.1914,  0.6905,
               -1.0859, -0.4433, -0.0255, 0.8198,  0.4571]]),
 'x': tensor([[-0.7205, -1.3996,  0.5528, -0.3059,  0.5384,
               -1.4976, -0.1480,  0.0841,0.3321,  0.5561]])}
>>> sample = p.sample({"y": sample_y, "a": sample_a}) # Redundant input ("a")
>>> print(sample) # input_var + var + "a" (redundant input)  # doctest: +SKIP
{'y': tensor([[ 1.3582, -1.1151, -0.8111,  1.0630,  1.1633,
                0.3855,  2.6324, -0.9357, -0.8649, -0.6015]]),
 'a': tensor([[-0.1874,  1.7958, -1.4084, -2.5646,  1.0868,
               -0.7523, -0.0852, -2.4222, -0.3914, -0.9755]]),
 'x': tensor([[-0.3272, -0.5222, -1.3659,  1.8386,  2.3204,
                0.3686,  0.6311, -1.1208, 0.3656, -0.6683]])}
get_log_prob(x_dict, **kwargs)[source]

Giving variables, this method returns values of log-pdf.

Parameters:
  • x_dict (dict) – Input variables.
  • sum_features (bool, defaults to True) – Whether the output is summed across some dimensions which are specified by feature_dims.
  • feature_dims (list or NoneType, defaults to None) – Set dimensions to sum across the output.
Returns:

log_prob – Values of log-probability density/mass function.

Return type:

torch.Tensor

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> sample_x = torch.randn(1, 10) # Psuedo data
>>> log_prob = p1.log_prob({"x": sample_x})
>>> print(log_prob) # doctest: +SKIP
tensor([-16.1153])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> log_prob = p2.log_prob({"x": sample_x, "y": sample_y})
>>> print(log_prob) # doctest: +SKIP
tensor([-21.5251])
sample_mean(x_dict={})[source]

Return the mean of the distribution.

Parameters:x_dict (dict, defaults to {}) – Parameters of this distribution.

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> mean = p1.sample_mean()
>>> print(mean)
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> mean = p2.sample_mean({"y": sample_y})
>>> print(mean) # doctest: +SKIP
tensor([[-0.2189, -1.0310, -0.1917, -0.3085,  1.5190, -0.9037,  1.2559,  0.1410,
          1.2810, -0.6681]])
sample_variance(x_dict={})[source]

Return the variance of the distribution.

Parameters:x_dict (dict, defaults to {}) – Parameters of this distribution.

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> var = p1.sample_variance()
>>> print(var)
tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> var = p2.sample_variance({"y": sample_y})
>>> print(var) # doctest: +SKIP
tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])
input_var

Input variables of this distribution. Normally, it has same values as cond_var.

Type:list
distribution_name

Name of this distribution class.

Type:str

MarginalizeVarDistribution

class pixyz.distributions.MarginalizeVarDistribution(p, marginalize_list)[source]

Bases: pixyz.distributions.distributions.Distribution

Marginalize variables in Distribution.

p(x) = \int p(x,z) dz

Examples

>>> a = DistributionBase(var=["x"], cond_var=["z"])
>>> b = DistributionBase(var=["y"], cond_var=["z"])
>>> p_multi = a * b
>>> print(p_multi)
Distribution:
  p(x,y|z) = p(x|z)p(y|z)
Network architecture:
  DistributionBase(
    name=p, distribution_name=,
    var=['y'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
  )
  DistributionBase(
    name=p, distribution_name=,
    var=['x'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
  )
>>> p_marg = MarginalizeVarDistribution(p_multi, ["y"])
>>> print(p_marg)
Distribution:
  p(x|z) = \int p(x|z)p(y|z)dy
Network architecture:
  DistributionBase(
    name=p, distribution_name=,
    var=['y'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
  )
  DistributionBase(
    name=p, distribution_name=,
    var=['x'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
  )
__init__(p, marginalize_list)[source]
Parameters:
  • p (pixyz.distributions.Distribution (not pixyz.distributions.DistributionBase)) – Distribution.
  • marginalize_list (list) – Variables to marginalize.
forward(*args, **kwargs)[source]

When this class is inherited by DNNs, this method should be overrided.

get_params(params_dict={})[source]

This method aims to get parameters of this distributions from constant parameters set in initialization and outputs of DNNs.

Parameters:params_dict (dict, defaults to {}) – Input parameters.
Returns:output_dict – Output parameters.
Return type:dict

Examples

>>> from pixyz.distributions import Normal
>>> dist_1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...                 features_shape=[1])
>>> print(dist_1)
Distribution:
  p(x)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([1])
    (loc): torch.Size([1, 1])
    (scale): torch.Size([1, 1])
  )
>>> dist_1.get_params()
{'loc': tensor([[0.]]), 'scale': tensor([[1.]])}
>>> dist_2 = Normal(loc=torch.tensor(0.), scale="z", cond_var=["z"], var=["x"])
>>> print(dist_2)
Distribution:
  p(x|z)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
    (loc): torch.Size([1])
  )
>>> dist_2.get_params({"z": torch.tensor(1.)})
{'scale': tensor(1.), 'loc': tensor([0.])}
sample(x_dict={}, batch_n=None, sample_shape=torch.Size([]), return_all=True, reparam=False)[source]

Sample variables of this distribution. If cond_var is not empty, you should set inputs as dict.

Parameters:
  • x_dict (torch.Tensor, list, or dict, defaults to {}) – Input variables.
  • sample_shape (list or NoneType, defaults to torch.Size()) – Shape of generating samples.
  • batch_n (int, defaults to None.) – Set batch size of parameters.
  • return_all (bool, defaults to True) – Choose whether the output contains input variables.
  • reparam (bool, defaults to False.) – Choose whether we sample variables with re-parameterized trick.
Returns:

output – Samples of this distribution.

Return type:

dict

Examples

>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...            features_shape=[10, 2])
>>> print(p)
Distribution:
  p(x)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([10, 2])
    (loc): torch.Size([1, 10, 2])
    (scale): torch.Size([1, 10, 2])
  )
>>> p.sample()["x"].shape  # (batch_n=1, features_shape)
torch.Size([1, 10, 2])
>>> p.sample(batch_n=20)["x"].shape  # (batch_n, features_shape)
torch.Size([20, 10, 2])
>>> p.sample(batch_n=20, sample_shape=[40, 30])["x"].shape  # (sample_shape, batch_n, features_shape)
torch.Size([40, 30, 20, 10, 2])
>>> # Conditional distribution
>>> p = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...            features_shape=[10])
>>> print(p)
Distribution:
  p(x|y)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=['y'], input_var=['y'], features_shape=torch.Size([10])
    (scale): torch.Size([1, 10])
  )
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> sample_a = torch.randn(1, 10) # Psuedo data
>>> sample = p.sample({"y": sample_y})
>>> print(sample) # input_var + var  # doctest: +SKIP
{'y': tensor([[-0.5182,  0.3484,  0.9042,  0.1914,  0.6905,
               -1.0859, -0.4433, -0.0255, 0.8198,  0.4571]]),
 'x': tensor([[-0.7205, -1.3996,  0.5528, -0.3059,  0.5384,
               -1.4976, -0.1480,  0.0841,0.3321,  0.5561]])}
>>> sample = p.sample({"y": sample_y, "a": sample_a}) # Redundant input ("a")
>>> print(sample) # input_var + var + "a" (redundant input)  # doctest: +SKIP
{'y': tensor([[ 1.3582, -1.1151, -0.8111,  1.0630,  1.1633,
                0.3855,  2.6324, -0.9357, -0.8649, -0.6015]]),
 'a': tensor([[-0.1874,  1.7958, -1.4084, -2.5646,  1.0868,
               -0.7523, -0.0852, -2.4222, -0.3914, -0.9755]]),
 'x': tensor([[-0.3272, -0.5222, -1.3659,  1.8386,  2.3204,
                0.3686,  0.6311, -1.1208, 0.3656, -0.6683]])}
sample_mean(x_dict={})[source]

Return the mean of the distribution.

Parameters:x_dict (dict, defaults to {}) – Parameters of this distribution.

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> mean = p1.sample_mean()
>>> print(mean)
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> mean = p2.sample_mean({"y": sample_y})
>>> print(mean) # doctest: +SKIP
tensor([[-0.2189, -1.0310, -0.1917, -0.3085,  1.5190, -0.9037,  1.2559,  0.1410,
          1.2810, -0.6681]])
sample_variance(x_dict={})[source]

Return the variance of the distribution.

Parameters:x_dict (dict, defaults to {}) – Parameters of this distribution.

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> var = p1.sample_variance()
>>> print(var)
tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> var = p2.sample_variance({"y": sample_y})
>>> print(var) # doctest: +SKIP
tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])
input_var

Input variables of this distribution. Normally, it has same values as cond_var.

Type:list
distribution_name

Name of this distribution class.

Type:str
prob_factorized_text

Return a formula of the factorized probability distribution.

Type:str

MultiplyDistribution

class pixyz.distributions.MultiplyDistribution(a, b)[source]

Bases: pixyz.distributions.distributions.Distribution

Multiply by given distributions, e.g, p(x,y|z) = p(x|z,y)p(y|z). In this class, it is checked if two distributions can be multiplied.

p(x|z)p(z|y) -> Valid

p(x|z)p(y|z) -> Valid

p(x|z)p(y|a) -> Valid

p(x|z)p(z|x) -> Invalid (recursive)

p(x|z)p(x|y) -> Invalid (conflict)

Examples

>>> a = DistributionBase(var=["x"], cond_var=["z"])
>>> b = DistributionBase(var=["z"], cond_var=["y"])
>>> p_multi = MultiplyDistribution(a, b)
>>> print(p_multi)
Distribution:
  p(x,z|y) = p(x|z)p(z|y)
Network architecture:
  DistributionBase(
    name=p, distribution_name=,
    var=['z'], cond_var=['y'], input_var=['y'], features_shape=torch.Size([])
  )
  DistributionBase(
    name=p, distribution_name=,
    var=['x'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
  )
>>> b = DistributionBase(var=["y"], cond_var=["z"])
>>> p_multi = MultiplyDistribution(a, b)
>>> print(p_multi)
Distribution:
  p(x,y|z) = p(x|z)p(y|z)
Network architecture:
  DistributionBase(
    name=p, distribution_name=,
    var=['y'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
  )
  DistributionBase(
    name=p, distribution_name=,
    var=['x'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
  )
>>> b = DistributionBase(var=["y"], cond_var=["a"])
>>> p_multi = MultiplyDistribution(a, b)
>>> print(p_multi)
Distribution:
  p(x,y|z,a) = p(x|z)p(y|a)
Network architecture:
  DistributionBase(
    name=p, distribution_name=,
    var=['y'], cond_var=['a'], input_var=['a'], features_shape=torch.Size([])
  )
  DistributionBase(
    name=p, distribution_name=,
    var=['x'], cond_var=['z'], input_var=['z'], features_shape=torch.Size([])
  )
__init__(a, b)[source]
Parameters:
  • a (pixyz.Distribution) – Distribution.
  • b (pixyz.Distribution) – Distribution.
input_var

Input variables of this distribution. Normally, it has same values as cond_var.

Type:list
prob_factorized_text

Return a formula of the factorized probability distribution.

Type:str
sample(x_dict={}, batch_n=None, return_all=True, reparam=False, **kwargs)[source]

Sample variables of this distribution. If cond_var is not empty, you should set inputs as dict.

Parameters:
  • x_dict (torch.Tensor, list, or dict, defaults to {}) – Input variables.
  • sample_shape (list or NoneType, defaults to torch.Size()) – Shape of generating samples.
  • batch_n (int, defaults to None.) – Set batch size of parameters.
  • return_all (bool, defaults to True) – Choose whether the output contains input variables.
  • reparam (bool, defaults to False.) – Choose whether we sample variables with re-parameterized trick.
Returns:

output – Samples of this distribution.

Return type:

dict

Examples

>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...            features_shape=[10, 2])
>>> print(p)
Distribution:
  p(x)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=[], input_var=[], features_shape=torch.Size([10, 2])
    (loc): torch.Size([1, 10, 2])
    (scale): torch.Size([1, 10, 2])
  )
>>> p.sample()["x"].shape  # (batch_n=1, features_shape)
torch.Size([1, 10, 2])
>>> p.sample(batch_n=20)["x"].shape  # (batch_n, features_shape)
torch.Size([20, 10, 2])
>>> p.sample(batch_n=20, sample_shape=[40, 30])["x"].shape  # (sample_shape, batch_n, features_shape)
torch.Size([40, 30, 20, 10, 2])
>>> # Conditional distribution
>>> p = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...            features_shape=[10])
>>> print(p)
Distribution:
  p(x|y)
Network architecture:
  Normal(
    name=p, distribution_name=Normal,
    var=['x'], cond_var=['y'], input_var=['y'], features_shape=torch.Size([10])
    (scale): torch.Size([1, 10])
  )
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> sample_a = torch.randn(1, 10) # Psuedo data
>>> sample = p.sample({"y": sample_y})
>>> print(sample) # input_var + var  # doctest: +SKIP
{'y': tensor([[-0.5182,  0.3484,  0.9042,  0.1914,  0.6905,
               -1.0859, -0.4433, -0.0255, 0.8198,  0.4571]]),
 'x': tensor([[-0.7205, -1.3996,  0.5528, -0.3059,  0.5384,
               -1.4976, -0.1480,  0.0841,0.3321,  0.5561]])}
>>> sample = p.sample({"y": sample_y, "a": sample_a}) # Redundant input ("a")
>>> print(sample) # input_var + var + "a" (redundant input)  # doctest: +SKIP
{'y': tensor([[ 1.3582, -1.1151, -0.8111,  1.0630,  1.1633,
                0.3855,  2.6324, -0.9357, -0.8649, -0.6015]]),
 'a': tensor([[-0.1874,  1.7958, -1.4084, -2.5646,  1.0868,
               -0.7523, -0.0852, -2.4222, -0.3914, -0.9755]]),
 'x': tensor([[-0.3272, -0.5222, -1.3659,  1.8386,  2.3204,
                0.3686,  0.6311, -1.1208, 0.3656, -0.6683]])}
get_log_prob(x_dict, sum_features=True, feature_dims=None)[source]

Giving variables, this method returns values of log-pdf.

Parameters:
  • x_dict (dict) – Input variables.
  • sum_features (bool, defaults to True) – Whether the output is summed across some dimensions which are specified by feature_dims.
  • feature_dims (list or NoneType, defaults to None) – Set dimensions to sum across the output.
Returns:

log_prob – Values of log-probability density/mass function.

Return type:

torch.Tensor

Examples

>>> import torch
>>> from pixyz.distributions import Normal
>>> # Marginal distribution
>>> p1 = Normal(loc=torch.tensor(0.), scale=torch.tensor(1.), var=["x"],
...             features_shape=[10], name="p1")
>>> sample_x = torch.randn(1, 10) # Psuedo data
>>> log_prob = p1.log_prob({"x": sample_x})
>>> print(log_prob) # doctest: +SKIP
tensor([-16.1153])
>>> # Conditional distribution
>>> p2 = Normal(loc="y", scale=torch.tensor(1.), var=["x"], cond_var=["y"],
...             features_shape=[10], name="p2")
>>> sample_y = torch.randn(1, 10) # Psuedo data
>>> log_prob = p2.log_prob({"x": sample_x, "y": sample_y})
>>> print(log_prob) # doctest: +SKIP
tensor([-21.5251])