Add functool.wraps() annotation to flax.nn.jit.

At the moment, all the jit names in a jaxpr show up as "jitted". functools.partial does not forward names. PiperOrigin-RevId: 648760671
google · Jul 2, 2024 · 37123d5 · 37123d5
1 parent 5b8265c
commit 37123d5
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/flax/linen/transforms.py b/flax/linen/transforms.py
@@ -608,8 +608,10 @@ def core_fn(
       return res
 
     core_fns = [
-      functools.partial(core_fn, prewrapped_fn, class_fn)
-      for prewrapped_fn, class_fn in zip(prewrapped_fns, class_fns)
+        functools.wraps(class_fn)(
+            functools.partial(core_fn, prewrapped_fn, class_fn)
+        )
+        for prewrapped_fn, class_fn in zip(prewrapped_fns, class_fns)
     ]
 
     # here we apply the given lifting transform to the scope-ingesting fn