Add custom derivative for scipy.special.hyp2f1

[mattbahr]

The scipy.special.hyp2f1 function was added here, but because of its use of lax.while_loop, the implementation does not work with autodiff. I think it would be beneficial to implement a custom derivative, similarly to how it's already been done for the scipy.special.hyp1f1 function.

jax/jax/_src/scipy/special.py

Lines 2642 to 2646 in d39f29c

	hyp1f1.defjvps(
	lambda a_dot, primal_out, a, b, x: _hyp1f1_a_derivative(a, b, x) * a_dot,
	lambda b_dot, primal_out, a, b, x: _hyp1f1_b_derivative(a, b, x) * b_dot,
	lambda x_dot, primal_out, a, b, x: _hyp1f1_x_derivative(a, b, x) * x_dot
	)

Based on this open issue with PyTensor, it seems that there's a need.

[jakevdp]

I agree in principle – the challenge lies in how to implement that derivative! There's probably a suitable series representation documented somewhere, but I'm not familiar with it.

[SuriyaaMM]

Found this from Tensorflow's Implementation
But even in Tensorflow's implementation they're skipping the gradients for a,b,c

def grad(dy):
      grad_z = a * b  * dy * hyp2f1_small_argument(
          a + 1., b + 1., c + 1., z) / c
      # We don't have an easily computable gradient with respect to parameters,
      # so ignore that for now.
      broadcast_shape = functools.reduce(
          ps.broadcast_shape,
          [ps.shape(x) for x in [a, b, c]])

      _, grad_z = tfp_math.fix_gradient_for_broadcasting(
          [tf.ones(broadcast_shape, dtype=z.dtype), z],
          [tf.ones_like(grad_z), grad_z])
      return None, None, None, grad_z

[SuriyaaMM]

can i work on this?
@jakevdp