Nathan Simpson

 Now neos works directly from this branch :)

Uploading some images for comparison -- this is a bit of a spanner in the works for our current project, so would really appreciate any pointers!  

Sagely advice @imurray, and sorry for not doing that sooner. I modified `examples/conditional_moons.ipynb` to have a two dimensional context feature, then worked it through and things worked as your message...

I suspected as much with `...` implying arbitrary shape, though I do think it's still helpful to add meaning to that shape within the output. Do you think it's enough...

> Hey @phinate, sorry for delay in response. Are you still interested in getting this merged in? > > > Do you think it's enough to say something like "`[context_size,...

> I assume this is also caused by this issue: > > ```python > from typing import List, Optional > from dataclasses import field, dataclass > > @dataclass > class...

Thanks both @mblondel & @fllinares for your replies, and the helpful information! I'm struggling a little with this because the suggestion of moving `param_for_grad` into the `to_minimize` call explicitly is...

Hi again all -- we followed @shoyer's suggestion of using `jax.closure_convert` using a function pretty much ripped from the `jax` docs: ```python3 def _minimize(objective_fn, lhood_pars, lr): converted_fn, aux_pars = jax.closure_convert(objective_fn,...

> Any follow up on this? Does your objective seem decomposable in the way I describe? Have just thought about this a bit -- I'm not 100% if this would...

Just an update on this: we've managed to get the jit working on our side with `closure_convert`, and we see the performance recover, so @shoyer got it right on that...