Composing Autoregressive and Fixed bijectors error

[matteobettini]

Issue Description

Hello, I am trying to compuse a bijector with AffineAutoregressive and Tanh. The flow crashes in this case when computing the log_prob.

Steps to Reproduce

dim = 2

base_dist = torch.distributions.Independent(
    torch.distributions.Normal(torch.zeros(dim), torch.ones(dim)), 1
)
bijectors = bij.Compose(
    [
        bij.AffineAutoregressive(
            params_fn=DenseAutoregressive(
                hidden_dims=(64, 64), nonlinearity=nn.ReLU
            )
        )
        for _ in range(1)
    ]
    + [bij.Tanh()]
)
flow = flowtorch.distributions.Flow(base_dist, bijectors)
flow.log_prob(torch.zeros(400, dim))

Traceback (most recent call last):
  File "/Users/Matteo/PycharmProjects/HetControl/het_control/flowtorch_test.py", line 73, in <module>
    flow.log_prob(torch.zeros(400, 2))
  File "/opt/homebrew/Caskroom/miniforge/base/envs/torchrl/lib/python3.9/site-packages/flowtorch/distributions/flow.py", line 114, in log_prob
    x = self.bijector.inverse(value, context)  # type: ignore
  File "/opt/homebrew/Caskroom/miniforge/base/envs/torchrl/lib/python3.9/site-packages/flowtorch/bijectors/compose.py", line 90, in inverse
    log_detJ = log_detJ + _log_detJ if log_detJ is not None else _log_detJ
RuntimeError: The size of tensor a (2) must match the size of tensor b (400) at non-singleton dimension 1

Expected Behavior

I expected to be able to chain compose the bijectors.

[matteobettini]

I would have posted this in https://github.com/stefanwebb/flowtorch but there is no issue tab there