Delayed param

[ordabayevy]

Addresses #533

Group coded with @fritzo @eb8680 @fehiepsi

[ordabayevy]

examples/minipyro.py is currently failing with jit. My guess is that jitting needs to be baked into funsor.adam.Adam, I will think more about this.

[ordabayevy]

Am I right that when using funsor.adam.Adam the function that needs to be jit traced is the loss function below (Subs funsor) @fritzo @eb8680 ? If yes, then it first needs to be converted to a function with positional arguments?

step_loss = loss(**{k: v[...] for k, v in params.items()}).data

[fritzo]

...failing with jit. My guess is that jitting needs to be baked into funsor.adam.Adam

I think you're right, but let's discuss. That's a little different from Pyro where jit is baked into ELBO subclasses.