Skip to content

Commit

Permalink
added docstring for Module.scope.path
Browse files Browse the repository at this point in the history
  • Loading branch information
chiamp committed May 8, 2024
1 parent b000c75 commit fdc893a
Showing 1 changed file with 26 additions and 0 deletions.
26 changes: 26 additions & 0 deletions flax/linen/module.py
Original file line number Diff line number Diff line change
Expand Up @@ -1552,6 +1552,32 @@ def _initialization_allowed(self):

@property
def path(self):
"""Get the path of this Module. Top-level root modules have an empty path ``()``.
Note that this method can only be used on bound modules that have a valid scope.
Example usage::
>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp
>>> class SubModel(nn.Module):
... @nn.compact
... def __call__(self, x):
... print(f'SubModel path: {self.path}')
... return x
>>> class Model(nn.Module):
... @nn.compact
... def __call__(self, x):
... print(f'Model path: {self.path}')
... return SubModel()(x)
>>> model = Model()
>>> variables = model.init(jax.random.key(0), jnp.ones((1, 2)))
Model path: ()
SubModel path: ('SubModel_0',)
"""

if self.scope is None:
raise ValueError("Can't access module paths on unbound modules.")

Expand Down

0 comments on commit fdc893a

Please sign in to comment.