In pytorch 1.14.0.dev20221128 when using vmap with greater_equal, I get the following warning:
UserWarning: There is a performance drop because we have not yet implemented the batching rule for aten::greater_equal.Scalar. Please file us an issue on GitHub so that we can prioritize its implementation. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/functorch/BatchedFallback.cpp:84.)
idx = torch.argmax(torch.greater_equal(torch.cumsum(input, dim=0), quantile).to(torch.int))
Alternatives
The alternative is to experience a performance drop. (The computation does complete correctly.)
🚀 The feature, motivation and pitch
In pytorch 1.14.0.dev20221128 when using
vmap
withgreater_equal
, I get the following warning:Alternatives
The alternative is to experience a performance drop. (The computation does complete correctly.)
Additional context
No response
cc @zou3519 @Chillee @samdow @soumith