Grad of constant?

[EngineerKhan]

I am trying a simple function:

def Func(x):
    if x>0:
        return x
    else:
        return 0

and when I try grad(Func)(a) where a is any non-positive number, it results in the following error:

TypeError: grad requires real-valued outputs (output dtype that is a sub-dtype of np.floating), but got int32. For differentiation of functions with integer outputs, use jax.vjp directly.

Shouldn't it simply return zero? Or am I missing something?

[jakevdp]

Try return 0.0 instead