Elementwise Composition gives AssertionError

[mlkrock]

Issue Description

Trying to compose a sigmoid with another bijector, but get an assertion error.

Steps to Reproduce

import torch
import flowtorch.bijectors as B
import flowtorch.distributions as D
import flowtorch.parameters as P
bijectors = B.Compose(bijectors  =[B.Sigmoid(), B.AffineAutoregressive()])
base_dist = torch.distributions.Independent(torch.distributions.Normal(torch.zeros(2), torch.ones(2)), 1)
flow = D.Flow(base_dist, bijectors)
opt = torch.optim.Adam(flow.parameters(), lr=5e-3)

gives

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/mitchell/Library/Python/3.9/lib/python/site-packages/torch/optim/adam.py", line 81, in __init__
    super(Adam, self).__init__(params, defaults)
  File "/Users/mitchell/Library/Python/3.9/lib/python/site-packages/torch/optim/optimizer.py", line 47, in __init__
    param_groups = list(params)
  File "/Users/mitchell/Library/Python/3.9/lib/python/site-packages/flowtorch/distributions/flow.py", line 50, in parameters
    for p in self.bijector.parameters():  # type: ignore
  File "/Users/mitchell/Library/Python/3.9/lib/python/site-packages/flowtorch/bijectors/compose.py", line 38, in parameters
    for param in b.parameters():  # type: ignore
  File "/Users/mitchell/Library/Python/3.9/lib/python/site-packages/flowtorch/bijectors/base.py", line 52, in parameters
    assert self._params_fn is not None
AssertionError

Expected Behavior

Sigmoid applied as elementwise composition after applying AffineAutoregressive(). If B.Sigmoid() is replaced with B.AffineAutoregressive() there is no error. Also got error if B.Exp() is used instead of B.Sigmoid(). Didn't try any others.

System Info

• PyTorch Version 1.11.0
• Python version 3.9.0

[stefanwebb]

Hi @mlkrock, thanks for reporting this bug! :)

It seems to have been inadvertently fixed with the latest commit on head. I have now done a new release for this commit, so if you upgrade from 0.6 to 0.7 it should be resolved:

pip install flowtorch==0.7

In general, I think it's best to install from Github source directly until the 1.0 release, and pull head as new commits are landed:

git clone https://github.com/facebookincubator/flowtorch.git
cd flowtorch
pip install -e .[dev]

[stefanwebb]

Let me know if this doesn't work and we can reopen the issue :)

[mlkrock]

Hi Stefan, Thanks, upgrading worked to fix the issue, but there is another one that appeared when I tried to optimize.

import torch
import flowtorch.bijectors as B
import flowtorch.distributions as D
import flowtorch.parameters as P

bijectors = B.Compose(bijectors  =[B.Sigmoid(), B.AffineAutoregressive()])
base_dist = torch.distributions.Independent(torch.distributions.Normal(torch.zeros(2), torch.ones(2)), 1)
flow = D.Flow(base_dist, bijectors)
opt = torch.optim.Adam(flow.parameters(), lr=5e-3)
target_dist = torch.distributions.Independent(torch.distributions.Normal(torch.zeros(2)+5, torch.ones(2)*0.5),1)
frame = 0

for idx in range(3001):
    opt.zero_grad()
    y = target_dist.sample((1000,))
    loss = -flow.log_prob(y).mean()
    if idx % 500 == 0:
        print('epoch', idx, 'loss', loss)
    loss.backward()
    opt.step()

gives the error

Traceback (most recent call last):
  File "<stdin>", line 4, in <module>
  File "/Users/mitchell/Library/Python/3.9/lib/python/site-packages/flowtorch/distributions/flow.py", line 125, in log_prob
    x = self.bijector.inverse(value, context)  # type: ignore
  File "/Users/mitchell/Library/Python/3.9/lib/python/site-packages/flowtorch/bijectors/compose.py", line 88, in inverse
    log_detJ = log_detJ + _log_detJ if log_detJ is not None else _log_detJ
RuntimeError: The size of tensor a (1000) must match the size of tensor b (2) at non-singleton dimension 1

Runs correctly if the Sigmoid is removed.

[stefanwebb]

Ah no, let me look into this bug...

[stefanwebb]

@mlkrock thanks for reporting this bug! I've been able to give a fix in #107.

The issue was that a previous refactoring introduced a bug where we were no longer taking into account the different event_dim's in a sequence of composed Bijectors.

Also, I just wanted to point out, is

bijectors = B.Compose(bijectors=[B.Sigmoid(), B.AffineAutoregressive()])

really what you mean to write? The Bijectors are applied in the order of the array, so it will pass the normal distribution through the sigmoid, and then through the affine autoregressive one.

Did you mean to apply the sigmoid last so it can transform the flow to lie on (0, 1)?

[mlkrock]

Hi @stefanwebb thanks for the fix, I tried and it is working nicely. Also thanks for the clarification, you are correct I wanted output to be (0,1).

[stefanwebb]

Glad I could be of assistance! :)

I'll close this issue now, feel free to submit again if you come across any bugs!