Removed private jax._src.api_util import and copied them into flax

Author: vfdev-5

flax/nnx/transforms/compilation.py

@@ -16,11 +16,12 @@
 
 import dataclasses
 import functools
+import inspect
+import operator
 import typing as tp
 
 import jax
 from jax.sharding import AbstractMesh, Mesh, PartitionSpec
-from jax._src import api_util  # We use fun_signature and resolve_argnums
 
 from flax.nnx import (
   extract,
@@ -390,17 +391,9 @@ def __init__(
       out_shardings,
     )
 
-    if isinstance(in_shardings, (list, tuple)):
+    if isinstance(in_shardings, (tuple, list)) and (static_argnums or static_argnames):
       # We should reintroduce None values into in_shardings corresponding to static arguments
-      fun_signature = api_util.fun_signature(fun)
-      _, _, static_argnums, _ = api_util.resolve_argnums(
-          fun,
-          fun_signature,
-          None,
-          None,
-          static_argnums,
-          static_argnames,
-      )
+      static_argnums = resolve_argnums(fun, static_argnums, static_argnames)
       in_shardings = list(in_shardings)
       for static_arg_index in sorted(static_argnums):
         in_shardings.insert(static_arg_index, None)
@@ -1051,3 +1044,113 @@ def shard_map_wrapper(*args, **kwargs):
   shard_map_wrapper.inner = shard_map_fn  # type: ignore
 
   return shard_map_wrapper  # type: ignore
+
+
+# We can't use private methods from jax._src.api_util
+# We copy the function: api_util.fun_signature
+def fun_signature(fun: tp.Callable) -> inspect.Signature | None:
+  try:
+    return inspect.signature(fun)
+  except (ValueError, TypeError):
+    return None
+
+# Adapted copy of private jax function from api_util: fun_signature
+def resolve_argnums(
+    fun: tp.Callable,
+    static_argnums: int | tp.Sequence[int] | None,
+    static_argnames: str | tp.Iterable[str] | None,
+) -> tuple[tuple[int, ...], tuple[str, ...], tuple[int, ...], tuple[str, ...]]:
+  def _ensure_index_tuple(x: tp.Any) -> tuple[int, ...]:
+    """Convert x to a tuple of indices."""
+    try:
+      return (operator.index(x),)
+    except TypeError:
+      return tuple(map(operator.index, x))
+
+  def _ensure_str(x: str) -> str:
+    if not isinstance(x, str):
+      raise TypeError(f"argument is not a string: {x}")
+    return x
+
+  def _ensure_str_tuple(x: str | tp.Iterable[str]) -> tuple[str, ...]:
+    """Convert x to a tuple of strings."""
+    if isinstance(x, str):
+      return (x,)
+    else:
+      return tuple(map(_ensure_str, x))
+
+  signature = fun_signature(fun)
+  if signature is None:
+    # Some built-in functions don't support signature.
+    # See: https://github.com/python/cpython/issues/73485
+    # In this case no validation is done
+    static_argnums = () if static_argnums is None else _ensure_index_tuple(
+        static_argnums)
+    static_argnames = () if static_argnames is None else _ensure_str_tuple(
+        static_argnames)
+  else:
+    # Infer argnums and argnames according to docstring
+    # If nums is None and names is not None, then nums are inferred from the
+    # names and vice-versa.
+    _POSITIONAL_OR_KEYWORD = inspect.Parameter.POSITIONAL_OR_KEYWORD
+    _POSITIONAL_ARGUMENTS = (
+      inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD
+    )
+    _KEYWORD_ARGUMENTS = (
+      inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY
+    )
+    _INVALID_KEYWORD_ARGUMENTS = (
+      inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.VAR_POSITIONAL
+    )
+
+    def infer_argnums_and_argnames(
+        sig: inspect.Signature,
+        argnums: int | tp.Iterable[int] | None,
+        argnames: str | tp.Iterable[str] | None,
+      ) -> tuple[tuple[int, ...], tuple[str, ...]]:
+      """Infer missing argnums and argnames for a function with inspect."""
+      if argnums is None and argnames is None:
+        return (), ()
+
+      if argnums is not None and argnames is not None:
+        argnums = _ensure_index_tuple(argnums)
+        argnames = _ensure_str_tuple(argnames)
+        return argnums, argnames
+
+      parameters = sig.parameters
+      if argnums is None:
+        assert argnames is not None
+        argnames = _ensure_str_tuple(argnames)
+        argnums = tuple(
+            i for i, (k, param) in enumerate(parameters.items())
+            if param.kind == _POSITIONAL_OR_KEYWORD and k in argnames
+        )
+      else:
+        argnums = _ensure_index_tuple(argnums)
+        argnames = tuple(
+            k for i, (k, param) in enumerate(parameters.items())
+            if param.kind == _POSITIONAL_OR_KEYWORD and i in argnums
+        )
+      return argnums, argnames
+
+    def _validate_argnums(sig: inspect.Signature, argnums: tuple[int, ...], argnums_name: str) -> None:
+      n_pos_args = 0
+      for param in sig.parameters.values():
+        if param.kind in _POSITIONAL_ARGUMENTS:
+          n_pos_args += 1
+
+        elif param.kind is inspect.Parameter.VAR_POSITIONAL:
+          # We can have any number of positional arguments
+          return
+
+      if argnums and (-min(argnums) > n_pos_args or max(argnums) >= n_pos_args):
+        raise ValueError(f"Jitted function has {argnums_name}={argnums}, "
+                        f"but only accepts {n_pos_args} positional arguments.")
+
+    static_argnums, static_argnames = infer_argnums_and_argnames(
+        signature, static_argnums, static_argnames)
+
+    # Validation
+    _validate_argnums(signature, static_argnums, "static_argnums")
+
+  return static_argnums