Minor documentation fixes for AxisMetadata.

PiperOrigin-RevId: 671753062
google · Sep 6, 2024 · cbf6c93 · cbf6c93
1 parent b494c3f
commit cbf6c93
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flax/core/meta.py b/flax/core/meta.py
@@ -58,13 +58,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.