jax

25183436 - use static_argnums in xla_computation (#2812)

use static_argnums in xla_computation (#2812) * use static_argnums in xla_computation fixes #1017 * add static_argnums to make_jaxpr * fix type error: handle int case