UserWarning: There is a performance drop because we have not yet implemented the batching rule for aten::greater.Scalar

[jonathan-macart]

🚀 The feature, motivation and pitch

In pytorch 2.0.1 (conda pytorch py3.8_cuda11.8_cudnn8.7.0_0) when using vmap with torch.greater, I get the following warning:

UserWarning: There is a performance drop because we have not yet implemented the batching rule for aten::greater.Scalar. Please file us an issue on GitHub so that we can prioritize its implementation. (Triggered internally at /opt/conda/conda-bld/pytorch_1682343998658/work/aten/src/ATen/functorch/BatchedFallback.cpp:82.)

MWE

import torch

def foo(x):
    return torch.greater(x, 1.0)

vmap_foo = torch.vmap(foo)
bar = torch.linspace(0, 10, 100)

assert torch.sum(vmap_foo(bar) ^ foo(bar)) == 0

Alternatives

The alternative is to experience a performance drop. (The computation does complete correctly.)

Additional context

Related: pytorch/functorch#1080

cc @zou3519 @Chillee @samdow @soumith @kshitij12345 @srossross

[kshitij12345]

This was fixed in #96744. If you try with nightly, you wouldn't see the warning there. Thanks!

[jonathan-macart]

Awesome; I can confirm the warning does not occur in the nightly. Thank you!