def test(a):
  b = np.arctan2(a,a)
  print(b)

  temp = np.array([0.,1.,0.])
  c = np.where(temp>0.5,0.,b)
  print(c)

  return np.sum(c)

aa = np.array([3.,0.,5.])
print(test(aa))
print(grad(test)(aa))