piecewise output same type as input

[hongwanliu]

Why does piecewise require the output to be the same as the input?

[jakevdp]

It's an implementation detail: the issue is that all functions in piecewise must return the same dtype, and casting to the input dtype was an easy way to make sure this condition was met.

Is this something that is blocking your use-case?

[hongwanliu]

Thanks. I can get around this for the case I was interested in, but it is a surprising requirement, and causes some inconvenience. I was trying to calculate the Riemann zeta function (which is very badly implemented in jax unfortunately: for negative, even integers it should be zero, but jax returns a value that is very much nonzero), and was hoping to pass integers and get floats as output.

[jakevdp]

Regarding the Zeta function, we just merged #17144 which should make things much better.

That aside, I can understand why the current casting behavior of jnp.piecewise is not ideal. I'll look into changing that.

[hongwanliu]

Awesome. Thanks!

[jakevdp]

#17186 should fix the piecewise issue.

[jakevdp]

Actually, it looks like the casting behavior you mention matches that of numpy:

In [1]: import numpy as np

In [2]: x = np.arange(10, dtype=int)

In [3]: np.piecewise(x, [x < 5, x > 5], [np.sin, np.cos]).dtype
Out[3]: dtype('int64')

Since the goal of jax.numpy functions is to provide the same API as numpy, but implemented in JAX, I think we should not merge #17186, but rather keep the behavior as-is. What do you think?

[hongwanliu]

Interesting. now that i’m aware of this for jax i’m ok either way, i just never needed piecewise in numpy whereas it’s much more important in jax.