Skip to content

Commit

Permalink
Minor documentation fixes for AxisMetadata.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 671753062
  • Loading branch information
Flax Team committed Sep 22, 2024
1 parent 268366f commit 491d930
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions flax/core/meta.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down

0 comments on commit 491d930

Please sign in to comment.