We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Improve tutorial on embedding networks
- make clear more early in the tutorial (and title) that this will be about embedding networks
- highlight more prominently pre-configured networks

--- 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/mixed_density_estimator.py`
Content:
```
1 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
2 # under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.
3 
4 from typing import Tuple
5 
6 import torch
7 from torch import Tensor
8 
9 from sbi.neural_nets.density_estimators import (
10     CategoricalMassEstimator,
11     DensityEstimator,
12 )
13 from sbi.neural_nets.density_estimators.nflows_flow import NFlowsFlow
14 from sbi.utils.sbiutils import match_theta_and_x_batch_shapes
15 from sbi.utils.torchutils import atleast_2d
16 
17 
18 class MixedDensityEstimator(DensityEstimator):
19     """Class performing Mixed Neural Likelihood Estimation.
20 
21     MNLE combines a Categorical net and a neural spline flow to model data with
22     mixed types, e.g., as they occur in decision-making models.
23     """
24 
25     def __init__(
26         self,
27         discrete_net: CategoricalMassEstimator,
28         continuous_net: NFlowsFlow,
29         condition_shape: torch.Size,
30         log_transform_x: bool = False,
31     ):
32         """Initialize class for combining density estimators for MNLE.
33 
34         Args:
35             discrete_net: neural net to model discrete part of the data.
36             continuous_net: neural net to model the continuous data.
37             log_transform_x: whether to transform the continous part of the data into
38                 logarithmic domain before training. This is helpful for bounded data, e.
39                 g.,for reaction times.
40         """
41         super(MixedDensityEstimator, self).__init__(
42             net=continuous_net, condition_shape=condition_shape
43         )
44 
45         self.discrete_net = discrete_net
46         self.continuous_net = continuous_net
47         self.log_transform_x = log_transform_x
48 
49     def forward(self, x: Tensor):
50         raise NotImplementedError(
51             """The forward method is not implemented for MNLE, use '.sample(...)' to
52             generate samples though a forward pass."""
53         )
54 
55     def sample(
56         self, context: Tensor, sample_shape: torch.Size, track_gradients: bool = False
57     ) -> Tensor:
58         """Return sample from mixed data distribution.
59 
60         Args:
61             context: parameters for which to generate samples.
62             sample_shape number of samples to generate.
63 
64         Returns:
65             Tensor: samples with shape (num_samples, num_data_dimensions)
66         """
67         assert (
68             context.shape[0] == 1
69         ), "Samples can be generated for a single context only."
70 
71         with torch.set_grad_enabled(track_gradients):
72             # Sample discrete data given parameters.
73             discrete_x = self.discrete_net.sample(
74                 sample_shape=sample_shape,
75                 context=context,
76             ).reshape(sample_shape[0], -1)
77 
78             # Sample continuous data condition on parameters and discrete data.
79             # Pass num_samples=1 because the choices in the context contains
80             # num_samples elements already.
81             continuous_x = self.continuous_net.sample(
82                 # repeat the single context to match number of sampled choices.
83                 # sample_shape[0] is the iid dimension.
84                 condition=torch.cat(
85                     (context.repeat(sample_shape[0], 1), discrete_x), dim=1
86                 ),
87                 sample_shape=sample_shape,
88             ).reshape(sample_shape[0], -1)
89 
90             # In case x was log-transformed, move them to linear space.
91             if self.log_transform_x:
92                 continuous_x = continuous_x.exp()
93 
94         return torch.cat((continuous_x, discrete_x), dim=1)
95 
96     def log_prob(self, x: Tensor, context: Tensor) -> Tensor:
97         """Return log-probability of samples under the learned MNLE.
98 
99         For a fixed data point x this returns the value of the likelihood function
100         evaluated at context, L(context, x=x).
101 
102         Alternatively, it can be interpreted as the log-prob of the density
103         p(x | context).
104 
105         It evaluates the separate density estimator for the discrete and continous part
106         of the data and then combines them into one evaluation.
107 
108         Args:
109             x: data (containing continuous and discrete data).
110             context: parameters for which to evaluate the likelihod function, or for
111                 which to condition p(x | context).
112 
113         Returns:
114             Tensor: log_prob of p(x | context).
115         """
116         assert (
117             x.shape[0] == context.shape[0]
118         ), "x and context must have same batch size."
119 
120         cont_x, disc_x = _separate_x(x)
121         dim_context = context.shape[0]
122 
123         disc_log_prob = self.discrete_net.log_prob(
124             input=disc_x, context=context
125         ).reshape(dim_context)
126 
127         cont_log_prob = self.continuous_net.log_prob(
128             # Transform to log-space if needed.
129             torch.log(cont_x) if self.log_transform_x else cont_x,
130             # Pass parameters and discrete x as context.
131             condition=torch.cat((context, disc_x), dim=1),
132         )
133 
134         # Combine into joint lp.
135         log_probs_combined = (disc_log_prob + cont_log_prob).reshape(dim_context)
136 
137         # Maybe add log abs det jacobian of RTs: log(1/x) = - log(x)
138         if self.log_transform_x:
139             log_probs_combined -= torch.log(cont_x).squeeze()
140 
141         return log_probs_combined
142 
143     def loss(self, x: Tensor, context: Tensor, **kwargs) -> Tensor:
144         return self.log_prob(x, context)
145 
146     def log_prob_iid(self, x: Tensor, context: Tensor) -> Tensor:
147         """Return log prob given a batch of iid x and a different batch of context.
148 
149         This is different from `.log_prob()` to enable speed ups in evaluation during
150         inference. The speed up is achieved by exploiting the fact that there are only
151         finite number of possible categories in the discrete part of the dat: one can
152         just calculate the log probs for each possible category (given the current batch
153         of context) and then copy those log probs into the entire batch of iid
154         categories.
155         For example, for the drift-diffusion model, there are only two choices, but
156         often 100s or 1000 trials. With this method a evaluation over trials then passes
157         a batch of `2 (one per choice) * num_contexts` into the NN, whereas the normal
158         `.log_prob()` would pass `1000 * num_contexts`.
159 
160         Args:
161             x: batch of iid data, data observed given the same underlying parameters or
162                 experimental conditions.
163             context: batch of parameters to be evaluated, i.e., each batch entry will be
164                 evaluated for the entire batch of iid x.
165 
166         Returns:
167             Tensor: log probs with shape (num_trials, num_parameters), i.e., the log
168                 prob for each context for each trial.
169         """
170 
171         context = atleast_2d(context)
172         x = atleast_2d(x)
173         batch_size = context.shape[0]
174         num_trials = x.shape[0]
175         context_repeated, x_repeated = match_theta_and_x_batch_shapes(context, x)
176         net_device = next(self.discrete_net.parameters()).device
177         assert (
178             net_device == x.device and x.device == context.device
179         ), f"device mismatch: net, x, context: \
180             {net_device}, {x.device}, {context.device}."
181 
182         x_cont_repeated, x_disc_repeated = _separate_x(x_repeated)
183         x_cont, x_disc = _separate_x(x)
184 
185         # repeat categories for parameters
186         repeated_categories = torch.repeat_interleave(
187             torch.arange(self.discrete_net.num_categories - 1), batch_size, dim=0
188         )
189         # repeat parameters for categories
190         repeated_context = context.repeat(self.discrete_net.num_categories - 1, 1)
191         log_prob_per_cat = torch.zeros(self.discrete_net.num_categories, batch_size).to(
192             net_device
193         )
194         log_prob_per_cat[:-1, :] = self.discrete_net.log_prob(
195             repeated_categories.to(net_device),
196             repeated_context.to(net_device),
197         ).reshape(-1, batch_size)
198         # infer the last category logprob from sum to one.
199         log_prob_per_cat[-1, :] = torch.log(1 - log_prob_per_cat[:-1, :].exp().sum(0))
200 
201         # fill in lps for each occurred category
202         log_probs_discrete = log_prob_per_cat[
203             x_disc.type_as(torch.zeros(1, dtype=torch.long)).squeeze()
204         ].reshape(-1)
205 
206         # Get repeat discrete data and context to match in batch shape for flow eval.
207         log_probs_cont = self.continuous_net.log_prob(
208             torch.log(x_cont_repeated) if self.log_transform_x else x_cont_repeated,
209             condition=torch.cat((context_repeated, x_disc_repeated), dim=1),
210         )
211 
212         # Combine into joint lp with first dim over trials.
213         log_probs_combined = (log_probs_discrete + log_probs_cont).reshape(
214             num_trials, batch_size
215         )
216 
217         # Maybe add log abs det jacobian of RTs: log(1/rt) = - log(rt)
218         if self.log_transform_x:
219             log_probs_combined -= torch.log(x_cont)
220 
221         # Return batch over trials as required by SBI potentials.
222         return log_probs_combined
223 
224 
225 def _separate_x(x: Tensor, num_discrete_columns: int = 1) -> Tuple[Tensor, Tensor]:
226     """Returns the continuous and discrete part of the given x.
227 
228     Assumes the discrete data to live in the last columns of x.
229     """
230 
231     assert x.ndim == 2, f"x must have two dimensions but has {x.ndim}."
232 
233     return x[:, :-num_discrete_columns], x[:, -num_discrete_columns:]
234
```
--- 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/mixed_density_estimator.py b/sbi/neural_nets/density_estimators/mixed_density_estimator.py
--- a/sbi/neural_nets/density_estimators/mixed_density_estimator.py
+++ b/sbi/neural_nets/density_estimators/mixed_density_estimator.py
@@ -164,8 +164,8 @@
                 evaluated for the entire batch of iid x.
 
         Returns:
-            Tensor: log probs with shape (num_trials, num_parameters), i.e., the log
-                prob for each context for each trial.
+            log probs with shape (num_trials, num_parameters), i.e., the log prob for
+            each context for each trial.
         """
 
         context = atleast_2d(context)

{"golden_diff": "diff --git a/sbi/neural_nets/density_estimators/mixed_density_estimator.py b/sbi/neural_nets/density_estimators/mixed_density_estimator.py\n--- a/sbi/neural_nets/density_estimators/mixed_density_estimator.py\n+++ b/sbi/neural_nets/density_estimators/mixed_density_estimator.py\n@@ -164,8 +164,8 @@\n                 evaluated for the entire batch of iid x.\n \n         Returns:\n-            Tensor: log probs with shape (num_trials, num_parameters), i.e., the log\n-                prob for each context for each trial.\n+            log probs with shape (num_trials, num_parameters), i.e., the log prob for\n+            each context for each trial.\n         \"\"\"\n \n         context = atleast_2d(context)\n", "issue": "Improve tutorial on embedding networks\n- make clear more early in the tutorial (and title) that this will be about embedding networks\r\n- highlight more prominently pre-configured networks\n", "before_files": [{"content": "# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed\n# under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.\n\nfrom typing import Tuple\n\nimport torch\nfrom torch import Tensor\n\nfrom sbi.neural_nets.density_estimators import (\n    CategoricalMassEstimator,\n    DensityEstimator,\n)\nfrom sbi.neural_nets.density_estimators.nflows_flow import NFlowsFlow\nfrom sbi.utils.sbiutils import match_theta_and_x_batch_shapes\nfrom sbi.utils.torchutils import atleast_2d\n\n\nclass MixedDensityEstimator(DensityEstimator):\n    \"\"\"Class performing Mixed Neural Likelihood Estimation.\n\n    MNLE combines a Categorical net and a neural spline flow to model data with\n    mixed types, e.g., as they occur in decision-making models.\n    \"\"\"\n\n    def __init__(\n        self,\n        discrete_net: CategoricalMassEstimator,\n        continuous_net: NFlowsFlow,\n        condition_shape: torch.Size,\n        log_transform_x: bool = False,\n    ):\n        \"\"\"Initialize class for combining density estimators for MNLE.\n\n        Args:\n            discrete_net: neural net to model discrete part of the data.\n            continuous_net: neural net to model the continuous data.\n            log_transform_x: whether to transform the continous part of the data into\n                logarithmic domain before training. This is helpful for bounded data, e.\n                g.,for reaction times.\n        \"\"\"\n        super(MixedDensityEstimator, self).__init__(\n            net=continuous_net, condition_shape=condition_shape\n        )\n\n        self.discrete_net = discrete_net\n        self.continuous_net = continuous_net\n        self.log_transform_x = log_transform_x\n\n    def forward(self, x: Tensor):\n        raise NotImplementedError(\n            \"\"\"The forward method is not implemented for MNLE, use '.sample(...)' to\n            generate samples though a forward pass.\"\"\"\n        )\n\n    def sample(\n        self, context: Tensor, sample_shape: torch.Size, track_gradients: bool = False\n    ) -> Tensor:\n        \"\"\"Return sample from mixed data distribution.\n\n        Args:\n            context: parameters for which to generate samples.\n            sample_shape number of samples to generate.\n\n        Returns:\n            Tensor: samples with shape (num_samples, num_data_dimensions)\n        \"\"\"\n        assert (\n            context.shape[0] == 1\n        ), \"Samples can be generated for a single context only.\"\n\n        with torch.set_grad_enabled(track_gradients):\n            # Sample discrete data given parameters.\n            discrete_x = self.discrete_net.sample(\n                sample_shape=sample_shape,\n                context=context,\n            ).reshape(sample_shape[0], -1)\n\n            # Sample continuous data condition on parameters and discrete data.\n            # Pass num_samples=1 because the choices in the context contains\n            # num_samples elements already.\n            continuous_x = self.continuous_net.sample(\n                # repeat the single context to match number of sampled choices.\n                # sample_shape[0] is the iid dimension.\n                condition=torch.cat(\n                    (context.repeat(sample_shape[0], 1), discrete_x), dim=1\n                ),\n                sample_shape=sample_shape,\n            ).reshape(sample_shape[0], -1)\n\n            # In case x was log-transformed, move them to linear space.\n            if self.log_transform_x:\n                continuous_x = continuous_x.exp()\n\n        return torch.cat((continuous_x, discrete_x), dim=1)\n\n    def log_prob(self, x: Tensor, context: Tensor) -> Tensor:\n        \"\"\"Return log-probability of samples under the learned MNLE.\n\n        For a fixed data point x this returns the value of the likelihood function\n        evaluated at context, L(context, x=x).\n\n        Alternatively, it can be interpreted as the log-prob of the density\n        p(x | context).\n\n        It evaluates the separate density estimator for the discrete and continous part\n        of the data and then combines them into one evaluation.\n\n        Args:\n            x: data (containing continuous and discrete data).\n            context: parameters for which to evaluate the likelihod function, or for\n                which to condition p(x | context).\n\n        Returns:\n            Tensor: log_prob of p(x | context).\n        \"\"\"\n        assert (\n            x.shape[0] == context.shape[0]\n        ), \"x and context must have same batch size.\"\n\n        cont_x, disc_x = _separate_x(x)\n        dim_context = context.shape[0]\n\n        disc_log_prob = self.discrete_net.log_prob(\n            input=disc_x, context=context\n        ).reshape(dim_context)\n\n        cont_log_prob = self.continuous_net.log_prob(\n            # Transform to log-space if needed.\n            torch.log(cont_x) if self.log_transform_x else cont_x,\n            # Pass parameters and discrete x as context.\n            condition=torch.cat((context, disc_x), dim=1),\n        )\n\n        # Combine into joint lp.\n        log_probs_combined = (disc_log_prob + cont_log_prob).reshape(dim_context)\n\n        # Maybe add log abs det jacobian of RTs: log(1/x) = - log(x)\n        if self.log_transform_x:\n            log_probs_combined -= torch.log(cont_x).squeeze()\n\n        return log_probs_combined\n\n    def loss(self, x: Tensor, context: Tensor, **kwargs) -> Tensor:\n        return self.log_prob(x, context)\n\n    def log_prob_iid(self, x: Tensor, context: Tensor) -> Tensor:\n        \"\"\"Return log prob given a batch of iid x and a different batch of context.\n\n        This is different from `.log_prob()` to enable speed ups in evaluation during\n        inference. The speed up is achieved by exploiting the fact that there are only\n        finite number of possible categories in the discrete part of the dat: one can\n        just calculate the log probs for each possible category (given the current batch\n        of context) and then copy those log probs into the entire batch of iid\n        categories.\n        For example, for the drift-diffusion model, there are only two choices, but\n        often 100s or 1000 trials. With this method a evaluation over trials then passes\n        a batch of `2 (one per choice) * num_contexts` into the NN, whereas the normal\n        `.log_prob()` would pass `1000 * num_contexts`.\n\n        Args:\n            x: batch of iid data, data observed given the same underlying parameters or\n                experimental conditions.\n            context: batch of parameters to be evaluated, i.e., each batch entry will be\n                evaluated for the entire batch of iid x.\n\n        Returns:\n            Tensor: log probs with shape (num_trials, num_parameters), i.e., the log\n                prob for each context for each trial.\n        \"\"\"\n\n        context = atleast_2d(context)\n        x = atleast_2d(x)\n        batch_size = context.shape[0]\n        num_trials = x.shape[0]\n        context_repeated, x_repeated = match_theta_and_x_batch_shapes(context, x)\n        net_device = next(self.discrete_net.parameters()).device\n        assert (\n            net_device == x.device and x.device == context.device\n        ), f\"device mismatch: net, x, context: \\\n            {net_device}, {x.device}, {context.device}.\"\n\n        x_cont_repeated, x_disc_repeated = _separate_x(x_repeated)\n        x_cont, x_disc = _separate_x(x)\n\n        # repeat categories for parameters\n        repeated_categories = torch.repeat_interleave(\n            torch.arange(self.discrete_net.num_categories - 1), batch_size, dim=0\n        )\n        # repeat parameters for categories\n        repeated_context = context.repeat(self.discrete_net.num_categories - 1, 1)\n        log_prob_per_cat = torch.zeros(self.discrete_net.num_categories, batch_size).to(\n            net_device\n        )\n        log_prob_per_cat[:-1, :] = self.discrete_net.log_prob(\n            repeated_categories.to(net_device),\n            repeated_context.to(net_device),\n        ).reshape(-1, batch_size)\n        # infer the last category logprob from sum to one.\n        log_prob_per_cat[-1, :] = torch.log(1 - log_prob_per_cat[:-1, :].exp().sum(0))\n\n        # fill in lps for each occurred category\n        log_probs_discrete = log_prob_per_cat[\n            x_disc.type_as(torch.zeros(1, dtype=torch.long)).squeeze()\n        ].reshape(-1)\n\n        # Get repeat discrete data and context to match in batch shape for flow eval.\n        log_probs_cont = self.continuous_net.log_prob(\n            torch.log(x_cont_repeated) if self.log_transform_x else x_cont_repeated,\n            condition=torch.cat((context_repeated, x_disc_repeated), dim=1),\n        )\n\n        # Combine into joint lp with first dim over trials.\n        log_probs_combined = (log_probs_discrete + log_probs_cont).reshape(\n            num_trials, batch_size\n        )\n\n        # Maybe add log abs det jacobian of RTs: log(1/rt) = - log(rt)\n        if self.log_transform_x:\n            log_probs_combined -= torch.log(x_cont)\n\n        # Return batch over trials as required by SBI potentials.\n        return log_probs_combined\n\n\ndef _separate_x(x: Tensor, num_discrete_columns: int = 1) -> Tuple[Tensor, Tensor]:\n    \"\"\"Returns the continuous and discrete part of the given x.\n\n    Assumes the discrete data to live in the last columns of x.\n    \"\"\"\n\n    assert x.ndim == 2, f\"x must have two dimensions but has {x.ndim}.\"\n\n    return x[:, :-num_discrete_columns], x[:, -num_discrete_columns:]\n", "path": "sbi/neural_nets/density_estimators/mixed_density_estimator.py"}], "after_files": [{"content": "# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed\n# under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.\n\nfrom typing import Tuple\n\nimport torch\nfrom torch import Tensor\n\nfrom sbi.neural_nets.density_estimators import (\n    CategoricalMassEstimator,\n    DensityEstimator,\n)\nfrom sbi.neural_nets.density_estimators.nflows_flow import NFlowsFlow\nfrom sbi.utils.sbiutils import match_theta_and_x_batch_shapes\nfrom sbi.utils.torchutils import atleast_2d\n\n\nclass MixedDensityEstimator(DensityEstimator):\n    \"\"\"Class performing Mixed Neural Likelihood Estimation.\n\n    MNLE combines a Categorical net and a neural spline flow to model data with\n    mixed types, e.g., as they occur in decision-making models.\n    \"\"\"\n\n    def __init__(\n        self,\n        discrete_net: CategoricalMassEstimator,\n        continuous_net: NFlowsFlow,\n        condition_shape: torch.Size,\n        log_transform_x: bool = False,\n    ):\n        \"\"\"Initialize class for combining density estimators for MNLE.\n\n        Args:\n            discrete_net: neural net to model discrete part of the data.\n            continuous_net: neural net to model the continuous data.\n            log_transform_x: whether to transform the continous part of the data into\n                logarithmic domain before training. This is helpful for bounded data, e.\n                g.,for reaction times.\n        \"\"\"\n        super(MixedDensityEstimator, self).__init__(\n            net=continuous_net, condition_shape=condition_shape\n        )\n\n        self.discrete_net = discrete_net\n        self.continuous_net = continuous_net\n        self.log_transform_x = log_transform_x\n\n    def forward(self, x: Tensor):\n        raise NotImplementedError(\n            \"\"\"The forward method is not implemented for MNLE, use '.sample(...)' to\n            generate samples though a forward pass.\"\"\"\n        )\n\n    def sample(\n        self, context: Tensor, sample_shape: torch.Size, track_gradients: bool = False\n    ) -> Tensor:\n        \"\"\"Return sample from mixed data distribution.\n\n        Args:\n            context: parameters for which to generate samples.\n            sample_shape number of samples to generate.\n\n        Returns:\n            Tensor: samples with shape (num_samples, num_data_dimensions)\n        \"\"\"\n        assert (\n            context.shape[0] == 1\n        ), \"Samples can be generated for a single context only.\"\n\n        with torch.set_grad_enabled(track_gradients):\n            # Sample discrete data given parameters.\n            discrete_x = self.discrete_net.sample(\n                sample_shape=sample_shape,\n                context=context,\n            ).reshape(sample_shape[0], -1)\n\n            # Sample continuous data condition on parameters and discrete data.\n            # Pass num_samples=1 because the choices in the context contains\n            # num_samples elements already.\n            continuous_x = self.continuous_net.sample(\n                # repeat the single context to match number of sampled choices.\n                # sample_shape[0] is the iid dimension.\n                condition=torch.cat(\n                    (context.repeat(sample_shape[0], 1), discrete_x), dim=1\n                ),\n                sample_shape=sample_shape,\n            ).reshape(sample_shape[0], -1)\n\n            # In case x was log-transformed, move them to linear space.\n            if self.log_transform_x:\n                continuous_x = continuous_x.exp()\n\n        return torch.cat((continuous_x, discrete_x), dim=1)\n\n    def log_prob(self, x: Tensor, context: Tensor) -> Tensor:\n        \"\"\"Return log-probability of samples under the learned MNLE.\n\n        For a fixed data point x this returns the value of the likelihood function\n        evaluated at context, L(context, x=x).\n\n        Alternatively, it can be interpreted as the log-prob of the density\n        p(x | context).\n\n        It evaluates the separate density estimator for the discrete and continous part\n        of the data and then combines them into one evaluation.\n\n        Args:\n            x: data (containing continuous and discrete data).\n            context: parameters for which to evaluate the likelihod function, or for\n                which to condition p(x | context).\n\n        Returns:\n            Tensor: log_prob of p(x | context).\n        \"\"\"\n        assert (\n            x.shape[0] == context.shape[0]\n        ), \"x and context must have same batch size.\"\n\n        cont_x, disc_x = _separate_x(x)\n        dim_context = context.shape[0]\n\n        disc_log_prob = self.discrete_net.log_prob(\n            input=disc_x, context=context\n        ).reshape(dim_context)\n\n        cont_log_prob = self.continuous_net.log_prob(\n            # Transform to log-space if needed.\n            torch.log(cont_x) if self.log_transform_x else cont_x,\n            # Pass parameters and discrete x as context.\n            condition=torch.cat((context, disc_x), dim=1),\n        )\n\n        # Combine into joint lp.\n        log_probs_combined = (disc_log_prob + cont_log_prob).reshape(dim_context)\n\n        # Maybe add log abs det jacobian of RTs: log(1/x) = - log(x)\n        if self.log_transform_x:\n            log_probs_combined -= torch.log(cont_x).squeeze()\n\n        return log_probs_combined\n\n    def loss(self, x: Tensor, context: Tensor, **kwargs) -> Tensor:\n        return self.log_prob(x, context)\n\n    def log_prob_iid(self, x: Tensor, context: Tensor) -> Tensor:\n        \"\"\"Return log prob given a batch of iid x and a different batch of context.\n\n        This is different from `.log_prob()` to enable speed ups in evaluation during\n        inference. The speed up is achieved by exploiting the fact that there are only\n        finite number of possible categories in the discrete part of the dat: one can\n        just calculate the log probs for each possible category (given the current batch\n        of context) and then copy those log probs into the entire batch of iid\n        categories.\n        For example, for the drift-diffusion model, there are only two choices, but\n        often 100s or 1000 trials. With this method a evaluation over trials then passes\n        a batch of `2 (one per choice) * num_contexts` into the NN, whereas the normal\n        `.log_prob()` would pass `1000 * num_contexts`.\n\n        Args:\n            x: batch of iid data, data observed given the same underlying parameters or\n                experimental conditions.\n            context: batch of parameters to be evaluated, i.e., each batch entry will be\n                evaluated for the entire batch of iid x.\n\n        Returns:\n            log probs with shape (num_trials, num_parameters), i.e., the log prob for\n            each context for each trial.\n        \"\"\"\n\n        context = atleast_2d(context)\n        x = atleast_2d(x)\n        batch_size = context.shape[0]\n        num_trials = x.shape[0]\n        context_repeated, x_repeated = match_theta_and_x_batch_shapes(context, x)\n        net_device = next(self.discrete_net.parameters()).device\n        assert (\n            net_device == x.device and x.device == context.device\n        ), f\"device mismatch: net, x, context: \\\n            {net_device}, {x.device}, {context.device}.\"\n\n        x_cont_repeated, x_disc_repeated = _separate_x(x_repeated)\n        x_cont, x_disc = _separate_x(x)\n\n        # repeat categories for parameters\n        repeated_categories = torch.repeat_interleave(\n            torch.arange(self.discrete_net.num_categories - 1), batch_size, dim=0\n        )\n        # repeat parameters for categories\n        repeated_context = context.repeat(self.discrete_net.num_categories - 1, 1)\n        log_prob_per_cat = torch.zeros(self.discrete_net.num_categories, batch_size).to(\n            net_device\n        )\n        log_prob_per_cat[:-1, :] = self.discrete_net.log_prob(\n            repeated_categories.to(net_device),\n            repeated_context.to(net_device),\n        ).reshape(-1, batch_size)\n        # infer the last category logprob from sum to one.\n        log_prob_per_cat[-1, :] = torch.log(1 - log_prob_per_cat[:-1, :].exp().sum(0))\n\n        # fill in lps for each occurred category\n        log_probs_discrete = log_prob_per_cat[\n            x_disc.type_as(torch.zeros(1, dtype=torch.long)).squeeze()\n        ].reshape(-1)\n\n        # Get repeat discrete data and context to match in batch shape for flow eval.\n        log_probs_cont = self.continuous_net.log_prob(\n            torch.log(x_cont_repeated) if self.log_transform_x else x_cont_repeated,\n            condition=torch.cat((context_repeated, x_disc_repeated), dim=1),\n        )\n\n        # Combine into joint lp with first dim over trials.\n        log_probs_combined = (log_probs_discrete + log_probs_cont).reshape(\n            num_trials, batch_size\n        )\n\n        # Maybe add log abs det jacobian of RTs: log(1/rt) = - log(rt)\n        if self.log_transform_x:\n            log_probs_combined -= torch.log(x_cont)\n\n        # Return batch over trials as required by SBI potentials.\n        return log_probs_combined\n\n\ndef _separate_x(x: Tensor, num_discrete_columns: int = 1) -> Tuple[Tensor, Tensor]:\n    \"\"\"Returns the continuous and discrete part of the given x.\n\n    Assumes the discrete data to live in the last columns of x.\n    \"\"\"\n\n    assert x.ndim == 2, f\"x must have two dimensions but has {x.ndim}.\"\n\n    return x[:, :-num_discrete_columns], x[:, -num_discrete_columns:]\n", "path": "sbi/neural_nets/density_estimators/mixed_density_estimator.py"}]}