Does there exist a at:where overload that can accept a float tensor instead of binary tensor? If it exists, then the mask allocation can also be eliminated (if needed).

@pytorchbot retest this please

@vadimkantorov Unfortunately not. This might be a good addition (and pretty easy to add), if there are other cases where we might need it. Maybe we should generalize this into some sort of arbitrary equality test against a floating point number? (CC @colesbury @cpuhrsch for more opinions).

I do have a question beyond what this patch does: when the exponent is close to zero (but not exactly zero), what happens to the gradients? Is it numerically stable? If it's not, it would be nice (though not strictly necessary) to fix that too.

@ezyang As the exponent tends to 0, the derivative is 0 for non-zero x, and tends to inf for x = 0. This works correctly in PyTorch:

>>> a
tensor([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.], requires_grad=True)
>>> a.pow(0.0001).sum().backward()
>>> a.grad
tensor([   inf, 0.0003, 0.0002, 0.0001, 0.0001, 0.0001, 0.0001, 0.0000, 0.0000,
        0.0000])
>>> a.pow(0.00001).sum().backward()
>>> a.grad
tensor([   inf, 0.0003, 0.0002, 0.0001, 0.0001, 0.0001, 0.0001, 0.0000, 0.0000,
        0.0000])
>>> a.pow(0.01).sum().backward()
>>> a.grad
tensor([   inf, 0.0103, 0.0052, 0.0035, 0.0026, 0.0021, 0.0017, 0.0015, 0.0013,
        0.0012])

@ezyang Thinking more about it, merging a comparison binary op and torch.where goes along the op fusion line (maybe it's a good case for the automatic fuser?). Comparison + torch.where is a frequent use-case I guess, but surfacing all comparison ops through where may bring unjustified api complexity.

The zero special case here would work because of a lucky coincidence: exponent needs to be compared to zero.

@vadimkantorov Certainly, "where" would be easy to support in the JIT fuser. I'd also be OK with a special case just for the zero test. Up to you guys!

I think for the purposes of this PR, the usage of where can remain as designed at the moment. Probably, when the semantics of where change, this part of the code can be revisited and modified accordingly. What do you think @ezyang ?

Works for me.