betainc batching rule is not quite correct

[fehiepsi]

This issue can be replicated using the following code

import jax
import jax.numpy as jnp
from jax.scipy.special import betainc

def f(x):
    return betainc(jnp.ones(3), 1., x)

assert f(jnp.ones(3)).shape == (3,)
assert jax.vmap(f)(jnp.ones((2, 3))).shape == (2, 3)
# TypeError: regularized_incomplete_beta got arrays of different rank: (2, 3), (2,), (2, 3).