jetstream-pytorch
jetstream-pytorch copied to clipboard
Support passing custom sampling function.
Users can pass in custom sampling function now. But it's not per request sampling because it's hard to track all the sampling functions in the jitted function. It applies to all the batches. The changes on the JetStream side is pending, not in yet.