Closed uri-granta closed 1 year ago
It might be a better idea to allow the shape to be specified / or have a flag like "relax_shapes" which does this.
Can you explain what you mean? Having the user specify the shape when creating the BatchReparametrizationSampler
(or when calling sample for the first time) feels like a major usability issue? And I though relax_shapes/reduce_retracing is about the arguments of compiled functions, not variable shapes?
Dynamic shape can have performance implications.
True, though worth remembering that it was already dynamically shaped on the previous release, so it's not a degradation.
Can you explain what you mean?
A few options:
Having the user specify the shape when creating the BatchReparametrizationSampler (or when calling sample for the first time) feels like a major usability issue?
I don't think either of these are usability problems really, because the default can always be no change. We already pass S, but you could default L and B to None and construct the new tensorshape from that. Equally, you could default the flag to be whichever way you wanted.
And I thought relax_shapes/reduce_retracing is about the arguments of compiled functions, not variable shapes?
Not sure what this means.
Can you explain what you mean?
A few options:
* pass in the shape of L, S and B * pass a flag which enables fixed_shape/relaxed_shapes, i.e. either do new or old behaviour behind a flag, so if shape is fixed it expects the same shape to always be used (but will infer it based on first use)
Ah, ok. I think adding a boolean option to disable the relaxed shapes does make sense. I'll do that in a subsequent PR as this one is blocking the docs build running successfully.
The recent commit to make reparameterization samplers more XLA friendly (#718) broke one of notebooks. Turns out that Async BO calls BatchReparametrizationSampler with different sized batches (due to pending points) meaning
._eps
needs to remain dynamically shaped in this case. Interestingly the JIT compilation test still passes.