Minor documentation fixes for AxisMetadata.

PiperOrigin-RevId: 677531963
google · Sep 22, 2024 · d584551 · d584551
1 parent 268366f
commit d584551
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flax/core/meta.py b/flax/core/meta.py
@@ -59,13 +59,13 @@ class AxisMetadata(Generic[A], metaclass=abc.ABCMeta):
   def unbox(self) -> A:
     """Returns the content of the AxisMetadata box.
 
-    Note that unlike ``meta.unbox`` the unbox call should recursively unbox
+    Note that unlike ``meta.unbox`` the unbox call should not recursively unbox
     metadata. It should simply return value that it wraps directly even
     if that value itself is an instance of AxisMetadata.
 
     In practise, AxisMetadata subclasses should be registered as PyTree nodes to
     support passing instances to JAX and Flax APIs. The leaves returned for this
-    note should correspond to the value returned by unbox.
+    node should correspond to the value returned by unbox.
 
     Returns:
       The unboxed value.