We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Density Estimator batched sample mixes up samples from different posteriors
**Describe the bug**
Given a batched observation, i.e., x1 and x2, the sampling method mixes up samples from different distributions.

**To Reproduce**
```python
import torch

from sbi import analysis as analysis
from sbi import utils as utils
from sbi.inference.base import infer

num_dim = 3
prior = utils.BoxUniform(low=-2 * torch.ones(num_dim), high=2 * torch.ones(num_dim))

def simulator(parameter_set):
    return 1.0 + parameter_set + torch.randn(parameter_set.shape) * 0.1

posterior = infer(simulator, prior, method="SNPE", num_simulations=200)
observation = torch.stack([torch.zeros(3), torch.ones(3)])
posterior_samples = posterior.posterior_estimator.sample((1000,), condition=observation)

# Outputs an multimodal distribution, but should be unimodal (mixes up samples from the two different x_os)
samples1 = posterior_samples[:,0].detach()
_ = analysis.pairplot([samples1], limits=[[-2, 2], [-2, 2], [-2, 2]], figsize=(6, 6))
```

**Additional context**

Likely a "reshaping" bug. 


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `sbi/neural_nets/density_estimators/nflows_flow.py`
Content:
```
1 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
2 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
3 
4 from typing import Tuple
5 
6 import torch
7 from pyknos.nflows.flows import Flow
8 from torch import Tensor, nn
9 
10 from sbi.neural_nets.density_estimators.base import DensityEstimator
11 from sbi.sbi_types import Shape
12 
13 
14 class NFlowsFlow(DensityEstimator):
15     r"""`nflows`- based normalizing flow density estimator.
16 
17     Flow type objects already have a .log_prob() and .sample() method, so here we just
18     wrap them and add the .loss() method.
19     """
20 
21     def __init__(
22         self, net: Flow, input_shape: torch.Size, condition_shape: torch.Size
23     ) -> None:
24         """Initialize density estimator which wraps flows from the `nflows` library.
25 
26         Args:
27             net: The raw `nflows` flow.
28             input_shape: Event shape of the input at which the density is being
29                 evaluated (and which is also the event_shape of samples).
30             condition_shape: Shape of the condition. If not provided, it will assume a
31                 1D input.
32         """
33         super().__init__(net, input_shape=input_shape, condition_shape=condition_shape)
34         # TODO: Remove as soon as DensityEstimator becomes abstract
35         self.net: Flow
36 
37     @property
38     def embedding_net(self) -> nn.Module:
39         r"""Return the embedding network."""
40         return self.net._embedding_net
41 
42     def inverse_transform(self, input: Tensor, condition: Tensor) -> Tensor:
43         r"""Return the inverse flow-transform of the inputs given a condition.
44 
45         The inverse transform is the transformation that maps the inputs back to the
46         base distribution (noise) space.
47 
48         Args:
49             input: Inputs to evaluate the inverse transform on of shape
50                     (*batch_shape1, input_size).
51             condition: Conditions of shape (*batch_shape2, *condition_shape).
52 
53         Raises:
54             RuntimeError: If batch_shape1 and batch_shape2 are not broadcastable.
55 
56         Returns:
57             noise: Transformed inputs.
58         """
59         self._check_condition_shape(condition)
60         condition_dims = len(self.condition_shape)
61 
62         # PyTorch's automatic broadcasting
63         batch_shape_in = input.shape[:-1]
64         batch_shape_cond = condition.shape[:-condition_dims]
65         batch_shape = torch.broadcast_shapes(batch_shape_in, batch_shape_cond)
66         # Expand the input and condition to the same batch shape
67         input = input.expand(batch_shape + (input.shape[-1],))
68         condition = condition.expand(batch_shape + self.condition_shape)
69         # Flatten required by nflows, but now both have the same batch shape
70         input = input.reshape(-1, input.shape[-1])
71         condition = condition.reshape(-1, *self.condition_shape)
72 
73         noise, _ = self.net._transorm(input, context=condition)
74         noise = noise.reshape(batch_shape)
75         return noise
76 
77     def log_prob(self, input: Tensor, condition: Tensor) -> Tensor:
78         r"""Return the log probabilities of the inputs given a condition or multiple
79         i.e. batched conditions.
80 
81         Args:
82             input: Inputs to evaluate the log probability on. Of shape
83                 `(sample_dim, batch_dim, *event_shape)`.
84             condition: Conditions of shape `(sample_dim, batch_dim, *event_shape)`.
85 
86         Raises:
87             AssertionError: If `input_batch_dim != condition_batch_dim`.
88 
89         Returns:
90             Sample-wise log probabilities, shape `(input_sample_dim, input_batch_dim)`.
91         """
92         input_sample_dim = input.shape[0]
93         input_batch_dim = input.shape[1]
94         condition_batch_dim = condition.shape[0]
95         condition_event_dims = len(condition.shape[1:])
96 
97         assert condition_batch_dim == input_batch_dim, (
98             f"Batch shape of condition {condition_batch_dim} and input "
99             f"{input_batch_dim} do not match."
100         )
101 
102         # Nflows needs to have a single batch dimension for condition and input.
103         input = input.reshape((input_batch_dim * input_sample_dim, -1))
104 
105         # Repeat the condition to match `input_batch_dim * input_sample_dim`.
106         ones_for_event_dims = (1,) * condition_event_dims  # Tuple of 1s, e.g. (1, 1, 1)
107         condition = condition.repeat(input_sample_dim, *ones_for_event_dims)
108 
109         log_probs = self.net.log_prob(input, context=condition)
110         return log_probs.reshape((input_sample_dim, input_batch_dim))
111 
112     def loss(self, input: Tensor, condition: Tensor) -> Tensor:
113         r"""Return the negative log-probability for training the density estimator.
114 
115         Args:
116             input: Inputs of shape `(batch_dim, *input_event_shape)`.
117             condition: Conditions of shape `(batch_dim, *condition_event_shape)`.
118 
119         Returns:
120             Negative log-probability of shape `(batch_dim,)`.
121         """
122         return -self.log_prob(input.unsqueeze(0), condition)[0]
123 
124     def sample(self, sample_shape: Shape, condition: Tensor) -> Tensor:
125         r"""Return samples from the density estimator.
126 
127         Args:
128             sample_shape: Shape of the samples to return.
129             condition: Conditions of shape `(sample_dim, batch_dim, *event_shape)`.
130 
131         Returns:
132             Samples of shape `(*sample_shape, condition_batch_dim)`.
133         """
134         condition_batch_dim = condition.shape[0]
135         num_samples = torch.Size(sample_shape).numel()
136 
137         samples = self.net.sample(num_samples, context=condition)
138 
139         return samples.reshape((
140             *sample_shape,
141             condition_batch_dim,
142             -1,
143         ))
144 
145     def sample_and_log_prob(
146         self, sample_shape: torch.Size, condition: Tensor, **kwargs
147     ) -> Tuple[Tensor, Tensor]:
148         r"""Return samples and their density from the density estimator.
149 
150         Args:
151             sample_shape: Shape of the samples to return.
152             condition: Conditions of shape (sample_dim, batch_dim, *event_shape).
153 
154         Returns:
155             Samples of shape `(*sample_shape, condition_batch_dim, *input_event_shape)`
156             and associated log probs of shape `(*sample_shape, condition_batch_dim)`.
157         """
158         condition_batch_dim = condition.shape[0]
159         num_samples = torch.Size(sample_shape).numel()
160 
161         samples, log_probs = self.net.sample_and_log_prob(
162             num_samples, context=condition
163         )
164         samples = samples.reshape((*sample_shape, condition_batch_dim, -1))
165         log_probs = log_probs.reshape((*sample_shape, -1))
166         return samples, log_probs
167
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/sbi/neural_nets/density_estimators/nflows_flow.py b/sbi/neural_nets/density_estimators/nflows_flow.py
--- a/sbi/neural_nets/density_estimators/nflows_flow.py
+++ b/sbi/neural_nets/density_estimators/nflows_flow.py
@@ -135,12 +135,8 @@
         num_samples = torch.Size(sample_shape).numel()
 
         samples = self.net.sample(num_samples, context=condition)
-
-        return samples.reshape((
-            *sample_shape,
-            condition_batch_dim,
-            -1,
-        ))
+        samples = samples.transpose(0, 1)
+        return samples.reshape((*sample_shape, condition_batch_dim, *self.input_shape))
 
     def sample_and_log_prob(
         self, sample_shape: torch.Size, condition: Tensor, **kwargs

{"golden_diff": "diff --git a/sbi/neural_nets/density_estimators/nflows_flow.py b/sbi/neural_nets/density_estimators/nflows_flow.py\n--- a/sbi/neural_nets/density_estimators/nflows_flow.py\n+++ b/sbi/neural_nets/density_estimators/nflows_flow.py\n@@ -135,12 +135,8 @@\n         num_samples = torch.Size(sample_shape).numel()\n \n         samples = self.net.sample(num_samples, context=condition)\n-\n-        return samples.reshape((\n-            *sample_shape,\n-            condition_batch_dim,\n-            -1,\n-        ))\n+        samples = samples.transpose(0, 1)\n+        return samples.reshape((*sample_shape, condition_batch_dim, *self.input_shape))\n \n     def sample_and_log_prob(\n         self, sample_shape: torch.Size, condition: Tensor, **kwargs\n", "issue": "Density Estimator batched sample mixes up samples from different posteriors\n**Describe the bug**\r\nGiven a batched observation, i.e., x1 and x2, the sampling method mixes up samples from different distributions.\r\n\r\n**To Reproduce**\r\n```python\r\nimport torch\r\n\r\nfrom sbi import analysis as analysis\r\nfrom sbi import utils as utils\r\nfrom sbi.inference.base import infer\r\n\r\nnum_dim = 3\r\nprior = utils.BoxUniform(low=-2 * torch.ones(num_dim), high=2 * torch.ones(num_dim))\r\n\r\ndef simulator(parameter_set):\r\n    return 1.0 + parameter_set + torch.randn(parameter_set.shape) * 0.1\r\n\r\nposterior = infer(simulator, prior, method=\"SNPE\", num_simulations=200)\r\nobservation = torch.stack([torch.zeros(3), torch.ones(3)])\r\nposterior_samples = posterior.posterior_estimator.sample((1000,), condition=observation)\r\n\r\n# Outputs an multimodal distribution, but should be unimodal (mixes up samples from the two different x_os)\r\nsamples1 = posterior_samples[:,0].detach()\r\n_ = analysis.pairplot([samples1], limits=[[-2, 2], [-2, 2], [-2, 2]], figsize=(6, 6))\r\n```\r\n\r\n**Additional context**\r\n\r\nLikely a \"reshaping\" bug. \r\n\n", "before_files": [{"content": "# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed\n# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>\n\nfrom typing import Tuple\n\nimport torch\nfrom pyknos.nflows.flows import Flow\nfrom torch import Tensor, nn\n\nfrom sbi.neural_nets.density_estimators.base import DensityEstimator\nfrom sbi.sbi_types import Shape\n\n\nclass NFlowsFlow(DensityEstimator):\n    r\"\"\"`nflows`- based normalizing flow density estimator.\n\n    Flow type objects already have a .log_prob() and .sample() method, so here we just\n    wrap them and add the .loss() method.\n    \"\"\"\n\n    def __init__(\n        self, net: Flow, input_shape: torch.Size, condition_shape: torch.Size\n    ) -> None:\n        \"\"\"Initialize density estimator which wraps flows from the `nflows` library.\n\n        Args:\n            net: The raw `nflows` flow.\n            input_shape: Event shape of the input at which the density is being\n                evaluated (and which is also the event_shape of samples).\n            condition_shape: Shape of the condition. If not provided, it will assume a\n                1D input.\n        \"\"\"\n        super().__init__(net, input_shape=input_shape, condition_shape=condition_shape)\n        # TODO: Remove as soon as DensityEstimator becomes abstract\n        self.net: Flow\n\n    @property\n    def embedding_net(self) -> nn.Module:\n        r\"\"\"Return the embedding network.\"\"\"\n        return self.net._embedding_net\n\n    def inverse_transform(self, input: Tensor, condition: Tensor) -> Tensor:\n        r\"\"\"Return the inverse flow-transform of the inputs given a condition.\n\n        The inverse transform is the transformation that maps the inputs back to the\n        base distribution (noise) space.\n\n        Args:\n            input: Inputs to evaluate the inverse transform on of shape\n                    (*batch_shape1, input_size).\n            condition: Conditions of shape (*batch_shape2, *condition_shape).\n\n        Raises:\n            RuntimeError: If batch_shape1 and batch_shape2 are not broadcastable.\n\n        Returns:\n            noise: Transformed inputs.\n        \"\"\"\n        self._check_condition_shape(condition)\n        condition_dims = len(self.condition_shape)\n\n        # PyTorch's automatic broadcasting\n        batch_shape_in = input.shape[:-1]\n        batch_shape_cond = condition.shape[:-condition_dims]\n        batch_shape = torch.broadcast_shapes(batch_shape_in, batch_shape_cond)\n        # Expand the input and condition to the same batch shape\n        input = input.expand(batch_shape + (input.shape[-1],))\n        condition = condition.expand(batch_shape + self.condition_shape)\n        # Flatten required by nflows, but now both have the same batch shape\n        input = input.reshape(-1, input.shape[-1])\n        condition = condition.reshape(-1, *self.condition_shape)\n\n        noise, _ = self.net._transorm(input, context=condition)\n        noise = noise.reshape(batch_shape)\n        return noise\n\n    def log_prob(self, input: Tensor, condition: Tensor) -> Tensor:\n        r\"\"\"Return the log probabilities of the inputs given a condition or multiple\n        i.e. batched conditions.\n\n        Args:\n            input: Inputs to evaluate the log probability on. Of shape\n                `(sample_dim, batch_dim, *event_shape)`.\n            condition: Conditions of shape `(sample_dim, batch_dim, *event_shape)`.\n\n        Raises:\n            AssertionError: If `input_batch_dim != condition_batch_dim`.\n\n        Returns:\n            Sample-wise log probabilities, shape `(input_sample_dim, input_batch_dim)`.\n        \"\"\"\n        input_sample_dim = input.shape[0]\n        input_batch_dim = input.shape[1]\n        condition_batch_dim = condition.shape[0]\n        condition_event_dims = len(condition.shape[1:])\n\n        assert condition_batch_dim == input_batch_dim, (\n            f\"Batch shape of condition {condition_batch_dim} and input \"\n            f\"{input_batch_dim} do not match.\"\n        )\n\n        # Nflows needs to have a single batch dimension for condition and input.\n        input = input.reshape((input_batch_dim * input_sample_dim, -1))\n\n        # Repeat the condition to match `input_batch_dim * input_sample_dim`.\n        ones_for_event_dims = (1,) * condition_event_dims  # Tuple of 1s, e.g. (1, 1, 1)\n        condition = condition.repeat(input_sample_dim, *ones_for_event_dims)\n\n        log_probs = self.net.log_prob(input, context=condition)\n        return log_probs.reshape((input_sample_dim, input_batch_dim))\n\n    def loss(self, input: Tensor, condition: Tensor) -> Tensor:\n        r\"\"\"Return the negative log-probability for training the density estimator.\n\n        Args:\n            input: Inputs of shape `(batch_dim, *input_event_shape)`.\n            condition: Conditions of shape `(batch_dim, *condition_event_shape)`.\n\n        Returns:\n            Negative log-probability of shape `(batch_dim,)`.\n        \"\"\"\n        return -self.log_prob(input.unsqueeze(0), condition)[0]\n\n    def sample(self, sample_shape: Shape, condition: Tensor) -> Tensor:\n        r\"\"\"Return samples from the density estimator.\n\n        Args:\n            sample_shape: Shape of the samples to return.\n            condition: Conditions of shape `(sample_dim, batch_dim, *event_shape)`.\n\n        Returns:\n            Samples of shape `(*sample_shape, condition_batch_dim)`.\n        \"\"\"\n        condition_batch_dim = condition.shape[0]\n        num_samples = torch.Size(sample_shape).numel()\n\n        samples = self.net.sample(num_samples, context=condition)\n\n        return samples.reshape((\n            *sample_shape,\n            condition_batch_dim,\n            -1,\n        ))\n\n    def sample_and_log_prob(\n        self, sample_shape: torch.Size, condition: Tensor, **kwargs\n    ) -> Tuple[Tensor, Tensor]:\n        r\"\"\"Return samples and their density from the density estimator.\n\n        Args:\n            sample_shape: Shape of the samples to return.\n            condition: Conditions of shape (sample_dim, batch_dim, *event_shape).\n\n        Returns:\n            Samples of shape `(*sample_shape, condition_batch_dim, *input_event_shape)`\n            and associated log probs of shape `(*sample_shape, condition_batch_dim)`.\n        \"\"\"\n        condition_batch_dim = condition.shape[0]\n        num_samples = torch.Size(sample_shape).numel()\n\n        samples, log_probs = self.net.sample_and_log_prob(\n            num_samples, context=condition\n        )\n        samples = samples.reshape((*sample_shape, condition_batch_dim, -1))\n        log_probs = log_probs.reshape((*sample_shape, -1))\n        return samples, log_probs\n", "path": "sbi/neural_nets/density_estimators/nflows_flow.py"}], "after_files": [{"content": "# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed\n# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>\n\nfrom typing import Tuple\n\nimport torch\nfrom pyknos.nflows.flows import Flow\nfrom torch import Tensor, nn\n\nfrom sbi.neural_nets.density_estimators.base import DensityEstimator\nfrom sbi.sbi_types import Shape\n\n\nclass NFlowsFlow(DensityEstimator):\n    r\"\"\"`nflows`- based normalizing flow density estimator.\n\n    Flow type objects already have a .log_prob() and .sample() method, so here we just\n    wrap them and add the .loss() method.\n    \"\"\"\n\n    def __init__(\n        self, net: Flow, input_shape: torch.Size, condition_shape: torch.Size\n    ) -> None:\n        \"\"\"Initialize density estimator which wraps flows from the `nflows` library.\n\n        Args:\n            net: The raw `nflows` flow.\n            input_shape: Event shape of the input at which the density is being\n                evaluated (and which is also the event_shape of samples).\n            condition_shape: Shape of the condition. If not provided, it will assume a\n                1D input.\n        \"\"\"\n        super().__init__(net, input_shape=input_shape, condition_shape=condition_shape)\n        # TODO: Remove as soon as DensityEstimator becomes abstract\n        self.net: Flow\n\n    @property\n    def embedding_net(self) -> nn.Module:\n        r\"\"\"Return the embedding network.\"\"\"\n        return self.net._embedding_net\n\n    def inverse_transform(self, input: Tensor, condition: Tensor) -> Tensor:\n        r\"\"\"Return the inverse flow-transform of the inputs given a condition.\n\n        The inverse transform is the transformation that maps the inputs back to the\n        base distribution (noise) space.\n\n        Args:\n            input: Inputs to evaluate the inverse transform on of shape\n                    (*batch_shape1, input_size).\n            condition: Conditions of shape (*batch_shape2, *condition_shape).\n\n        Raises:\n            RuntimeError: If batch_shape1 and batch_shape2 are not broadcastable.\n\n        Returns:\n            noise: Transformed inputs.\n        \"\"\"\n        self._check_condition_shape(condition)\n        condition_dims = len(self.condition_shape)\n\n        # PyTorch's automatic broadcasting\n        batch_shape_in = input.shape[:-1]\n        batch_shape_cond = condition.shape[:-condition_dims]\n        batch_shape = torch.broadcast_shapes(batch_shape_in, batch_shape_cond)\n        # Expand the input and condition to the same batch shape\n        input = input.expand(batch_shape + (input.shape[-1],))\n        condition = condition.expand(batch_shape + self.condition_shape)\n        # Flatten required by nflows, but now both have the same batch shape\n        input = input.reshape(-1, input.shape[-1])\n        condition = condition.reshape(-1, *self.condition_shape)\n\n        noise, _ = self.net._transorm(input, context=condition)\n        noise = noise.reshape(batch_shape)\n        return noise\n\n    def log_prob(self, input: Tensor, condition: Tensor) -> Tensor:\n        r\"\"\"Return the log probabilities of the inputs given a condition or multiple\n        i.e. batched conditions.\n\n        Args:\n            input: Inputs to evaluate the log probability on. Of shape\n                `(sample_dim, batch_dim, *event_shape)`.\n            condition: Conditions of shape `(sample_dim, batch_dim, *event_shape)`.\n\n        Raises:\n            AssertionError: If `input_batch_dim != condition_batch_dim`.\n\n        Returns:\n            Sample-wise log probabilities, shape `(input_sample_dim, input_batch_dim)`.\n        \"\"\"\n        input_sample_dim = input.shape[0]\n        input_batch_dim = input.shape[1]\n        condition_batch_dim = condition.shape[0]\n        condition_event_dims = len(condition.shape[1:])\n\n        assert condition_batch_dim == input_batch_dim, (\n            f\"Batch shape of condition {condition_batch_dim} and input \"\n            f\"{input_batch_dim} do not match.\"\n        )\n\n        # Nflows needs to have a single batch dimension for condition and input.\n        input = input.reshape((input_batch_dim * input_sample_dim, -1))\n\n        # Repeat the condition to match `input_batch_dim * input_sample_dim`.\n        ones_for_event_dims = (1,) * condition_event_dims  # Tuple of 1s, e.g. (1, 1, 1)\n        condition = condition.repeat(input_sample_dim, *ones_for_event_dims)\n\n        log_probs = self.net.log_prob(input, context=condition)\n        return log_probs.reshape((input_sample_dim, input_batch_dim))\n\n    def loss(self, input: Tensor, condition: Tensor) -> Tensor:\n        r\"\"\"Return the negative log-probability for training the density estimator.\n\n        Args:\n            input: Inputs of shape `(batch_dim, *input_event_shape)`.\n            condition: Conditions of shape `(batch_dim, *condition_event_shape)`.\n\n        Returns:\n            Negative log-probability of shape `(batch_dim,)`.\n        \"\"\"\n        return -self.log_prob(input.unsqueeze(0), condition)[0]\n\n    def sample(self, sample_shape: Shape, condition: Tensor) -> Tensor:\n        r\"\"\"Return samples from the density estimator.\n\n        Args:\n            sample_shape: Shape of the samples to return.\n            condition: Conditions of shape `(sample_dim, batch_dim, *event_shape)`.\n\n        Returns:\n            Samples of shape `(*sample_shape, condition_batch_dim)`.\n        \"\"\"\n        condition_batch_dim = condition.shape[0]\n        num_samples = torch.Size(sample_shape).numel()\n\n        samples = self.net.sample(num_samples, context=condition)\n        samples = samples.transpose(0, 1)\n        return samples.reshape((*sample_shape, condition_batch_dim, *self.input_shape))\n\n    def sample_and_log_prob(\n        self, sample_shape: torch.Size, condition: Tensor, **kwargs\n    ) -> Tuple[Tensor, Tensor]:\n        r\"\"\"Return samples and their density from the density estimator.\n\n        Args:\n            sample_shape: Shape of the samples to return.\n            condition: Conditions of shape (sample_dim, batch_dim, *event_shape).\n\n        Returns:\n            Samples of shape `(*sample_shape, condition_batch_dim, *input_event_shape)`\n            and associated log probs of shape `(*sample_shape, condition_batch_dim)`.\n        \"\"\"\n        condition_batch_dim = condition.shape[0]\n        num_samples = torch.Size(sample_shape).numel()\n\n        samples, log_probs = self.net.sample_and_log_prob(\n            num_samples, context=condition\n        )\n        samples = samples.reshape((*sample_shape, condition_batch_dim, -1))\n        log_probs = log_probs.reshape((*sample_shape, -1))\n        return samples, log_probs\n", "path": "sbi/neural_nets/density_estimators/nflows_flow.py"}]}