Bug when raising zero to powers?

[yairchu]

Consider these two seemingly equivalent functions:

def fa(x):
    return x ** 1.5

def fb(x):
    return x * x ** 0.5

We can see that one of them is differentiated ok, while the other produces warnings and nans at zero:

>>> [ga, gb] = map(autograd.grad, [fa, fb])
>>> print(ga(2.), ga(0.))
2.121320343559643 0.0
>>> print(gb(2.), gb(0.))
.../autograd/numpy/numpy_vjps.py:59: RuntimeWarning: divide by zero encountered in power
  lambda ans, x, y : unbroadcast_f(x, lambda g: g * y * x ** anp.where(y, y - 1, 1.)),
.../autograd/numpy/numpy_vjps.py:59: RuntimeWarning: invalid value encountered in double_scalars
  lambda ans, x, y : unbroadcast_f(x, lambda g: g * y * x ** anp.where(y, y - 1, 1.)),
2.121320343559643 nan

I'm not sure whether this is a proper fix, but if in the defvjp call of anp.power we change anp.where(y, y - 1, 1.) to anp.where(x, anp.where(y, y - 1, 1.), 1.) then gb(0.) does produce the same result as ga(0.)

I concede that 0 ** -0.5 and ZeroDivisionError in general is a delicate topic. But my fix still seems consistent to me. 🤷‍♂️