Wrong static_argnums in negative_binomial in release v0.3.2

[mhliu0001]

When I downloaded the latest release of appletree and tried to import, I encountered an error:

Python 3.10.14 (main, Mar 21 2024, 16:24:04) [GCC 11.2.0]
Type 'copyright', 'credits' or 'license' for more information
IPython 8.24.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: import appletree
Using Normal as an approximation of Binomial
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[1], line 1
----> 1 import appletree

File ~/anaconda3/envs/xenon_apt/lib/python3.10/site-packages/appletree/__init__.py:18
     14 from .interpolation import *
     16 from .config import *
---> 18 from .parameter import *
     20 from .randgen import *
     22 from .share import *

File ~/anaconda3/envs/xenon_apt/lib/python3.10/site-packages/appletree/parameter.py:6
      2 import json
      4 import numpy as np
----> 6 from appletree.randgen import TwoHalfNorm
      7 from appletree.utils import errors_to_two_half_norm_sigmas
     10 class Parameter:

File ~/anaconda3/envs/xenon_apt/lib/python3.10/site-packages/appletree/randgen.py:281
    275         ret = vmap(lambda *x: dispatch(*x))(seed, p, n)
    276     return key, jnp.reshape(ret, shape)
    279 @export
    280 @partial(jit, static_argnums=(3, 4))
--> 281 def negative_binomial(key, p, n, shape=()):
    282     """Negative binomial distribution random sampler. Using Gamma–Poisson mixture.
    283 
    284     Args:
   (...)
    297 
    298     """
    300     key, lam = gamma(key, n, p / (1 - p), shape)

File ~/anaconda3/envs/xenon_apt/lib/python3.10/site-packages/jax/_src/api.py:281, in jit(fun, in_shardings, out_shardings, static_argnums, static_argnames, donate_argnums, donate_argnames, keep_unused, device, backend, inline, abstracted_axes)
    142 def jit(
    143   fun: Callable,
    144   in_shardings=sharding_impls.UNSPECIFIED,
   (...)
    154   abstracted_axes: Any | None = None,
    155 ) -> pjit.JitWrapped:
    156   """Sets up ``fun`` for just-in-time compilation with XLA.
    157 
    158   Args:
   (...)
    279     Array([   0,    1,  256, 6561], dtype=int32)
    280   """
--> 281   return pjit.make_jit(
    282         fun, in_shardings, out_shardings, donate_argnums, donate_argnames,
    283         static_argnums, static_argnames, device, backend, abstracted_axes,
    284         keep_unused, inline, use_resource_env=False)

File ~/anaconda3/envs/xenon_apt/lib/python3.10/site-packages/jax/_src/pjit.py:500, in make_jit(fun, in_shardings, out_shardings, donate_argnums, donate_argnames, static_argnums, static_argnames, device, backend, abstracted_axes, keep_unused, inline, use_resource_env)
    491 def make_jit(fun: Callable, in_shardings: Any, out_shardings: Any,
    492              donate_argnums: int | Sequence[int] | None,
    493              donate_argnames: str | Iterable[str] | None,
   (...)
    497              abstracted_axes: Any | None, keep_unused: bool,
    498              inline: bool, use_resource_env: bool) -> Any:
    499   """jit() and pjit() are thin wrappers around this function."""
--> 500   jit_info = _parse_jit_arguments(
    501         fun, in_shardings, out_shardings, donate_argnums, donate_argnames,
    502         static_argnums, static_argnames, device, backend, abstracted_axes,
    503         keep_unused, inline, use_resource_env)
    504   return _make_jit_wrapper(jit_info)

File ~/anaconda3/envs/xenon_apt/lib/python3.10/site-packages/jax/_src/pjit.py:420, in _parse_jit_arguments(fun, in_shardings, out_shardings, donate_argnums, donate_argnames, static_argnums, static_argnames, device, backend, abstracted_axes, keep_unused, inline, use_resource_env)
    417 fun_sourceinfo = api_util.fun_sourceinfo(fun)
    418 fun_signature = api_util.fun_signature(fun)
--> 420 donate_argnums, donate_argnames, static_argnums, static_argnames = resolve_argnums(
    421     fun, fun_signature, donate_argnums, donate_argnames, static_argnums,
    422     static_argnames)
    424 has_explicit_sharding = _pjit_explicit_sharding(
    425     in_shardings, out_shardings, device, backend)
    427 return PjitInfo(
    428       fun=fun,
    429       fun_sourceinfo=fun_sourceinfo,
   (...)
    445       has_explicit_sharding=has_explicit_sharding,
    446       use_resource_env=use_resource_env)

File ~/anaconda3/envs/xenon_apt/lib/python3.10/site-packages/jax/_src/api_util.py:540, in resolve_argnums(fun, signature, donate_argnums, donate_argnames, static_argnums, static_argnames)
    536 donate_argnums, donate_argnames = infer_argnums_and_argnames(
    537     signature, donate_argnums, donate_argnames)
    539 # Validation
--> 540 _validate_argnums(signature, static_argnums, "static_argnums")
    541 _validate_argnames(signature, static_argnames, "static_argnames")
    542 _validate_argnums(signature, donate_argnums, "donate_argnums")

File ~/anaconda3/envs/xenon_apt/lib/python3.10/site-packages/jax/_src/api_util.py:170, in _validate_argnums(sig, argnums, argnums_name)
    167     return
    169 if argnums and (-min(argnums) > n_pos_args or max(argnums) >= n_pos_args):
--> 170   raise ValueError(f"Jitted function has {argnums_name}={argnums}, "
    171                    f"but only accepts {n_pos_args} positional arguments.")

ValueError: Jitted function has static_argnums=(3, 4), but only accepts 4 positional arguments.

This was introduced in PR #145 and fixed in PR #148. We might need to release one more time.