I've made some progress now towards this PR, and would like some feedback before I continue.

BasePotential can either "allow_iid" or not.

Givenbatch_dim_theta!=batch_dim_x, we need to decide how to interpret how to evaluatepotential(x,theta). We could return(batch_dim_x,batch_dim_theta) potentials (i.e. every combination), but I am worried this can add a lot of computational overhead, especially when sampling. Instead, the current implementation I suggest that we assume thatbatch_dim_theta is a multiple ofbatch_dim_x (i.e. for sampling, we have n chains in theta for eachx). In this case we expand the batch dim ofx tobatch_theta, and match whichx goes to whichtheta. If we are happy with this approach, I'll go ahead and apply this also to the MCMCinit_strategy, etc., and make sure this is consistent with other calls.

Remove warning for batched x and default to batched evaluation
Not sure if we want batched evaluation as the default. I think it's easier to do batched evaluation whensample_batched orlog_prob_batched is called, and otherwise assume iid (and warn if batch dim >1 as before).

Great, it looks good. I like that the choice on iid or not can now be made at theset_x method which makes a lot of sense.

I would also opt for your suggested option. The question arises because we squeeze the batch_shape into a single dimension, right? For "PyTorch" broadcasting, one would expect something like (1,batch_x_dim, x_dim) and (batch_theta_dim, betach_x_dim, theta_dim) -> (batch_x_dim, batch_theta_dim), so by squeezing the xs, thetas into 2d one would always get a dimension that is a multiple of batch_x_dim (otherwise it cannot be represented by a fixed size tensor).

For (1,batch_x_dim,x_dim) and (batch_theta_dim, 1, theta_dim), PyTorch broadcasting semantics would compute all combinations. Unfortunately, after squeezing, these distinctions between cases can no longer be fully preserved.

Great effort, thanks a lot for tacking this 👏

I do have a couple of comments and questions. Happy to discuss in person if needed.

Thanks for the review! I implemented your suggestions.

An additional point - Forposterior_based_potential, indeed we should not allow foriid_x, as this is handled byPermutationInvariantNetwork. Instead, we now always treatx batches as not iid. If the user tries to setpotential.set_x(x,x_is_iid=True) with aPosteriorBasedPotential, we raise an error stating this. I added a few test cases inembedding_net_test.py::test_embedding_api_with_multiple_trials to test whether batches ofx are interpreted correctly when we use aPermutationInvariantNetwork.

closes#990
closes#944