generative_function is broken for models with non-optional args

Given the model:

@mcx.model
def model(x):
    a <~ dist.Normal(0, 1)
    return a

This code: https://github.com/rlouf/mcx/blob/master/mcx/model.py#L218-L220

Will generate a function with a signature model_sample_posterior_predictive(rng_key, x, a)

This function will then be called here: https://github.com/rlouf/mcx/blob/master/mcx/model.py#L226-L228 passing *self.trace.chain_values(self.chain_id) as x.

The fix is to modify the generative_function.__call__ method to

    def __call__(self, rng_key, *args, **kwargs) -> jnp.DeviceArray:
        """Call the model as a generative function."""

        posterior_kwd = {k: self.trace[k].reshape(-1) for k in self.trace.keys()}
        return self.call_fn(
            rng_key, *args, **kwargs, **posterior_kwd,
        )

[elanmart]

@rlouf thanks a lot! In the meantime I’ve just patched my local version of mcx so I can carry on without issues. For #90 I also have a small workaround, so it’s all fine for now 😃

[rlouf]

(again, thank you for raising an issue!)