Support JAX for lambdify

[jjaraalm]

JAX is very useful and it would be nice is sympy supported lambdifying expressions using it. I can see a couple different features

• Swap numpy and scipy calls for jax.numpy and jax.scipy when available.
• Replace derivatives of scalar valued functions with JAX derivatives so that e.g., Derivative(f(x,y), (x,2)) maps to jax.grad(jax.grad(f, 0), 0)
• Jacobian or Jacobian-vector products of Array. This seems less straightforward or meaningful since each array element is treated separately and not required to have the same arguments.

[so-rose]

This works:

• https://docs.sympy.org/latest/modules/numeric-computation.html#lambdify
• https://docs.sympy.org/latest/tutorials/intro-tutorial/basic_operations.html#lambdify

Array broadcasting, grad, etc. is all fine - Just set the mode='jax'.