diff --git a/flax/cursor.py b/flax/cursor.py index 334e23456..60a98ddb4 100644 --- a/flax/cursor.py +++ b/flax/cursor.py @@ -396,7 +396,7 @@ def find(self, cond_fn: Callable[[str, Any], bool]) -> 'Cursor[A]': - If the ``cond_fn`` evaluates to True at a particular key path, this method will not recurse any further down that branch; i.e. this method will find and return the "earliest" child node that fulfills the condition in ``cond_fn`` in a particular key path - - ``.find`` WILL NOT search the the value at the top-most level of the pytree (i.e. the root + - ``.find`` WILL NOT search the value at the top-most level of the pytree (i.e. the root node). The ``cond_fn`` will be evaluated recursively, starting at the root node's children. Example:: @@ -484,7 +484,7 @@ def find_all( - If the ``cond_fn`` evaluates to True at a particular key path, this method will not recurse any further down that branch; i.e. this method will find and return the "earliest" child nodes that fulfill the condition in ``cond_fn`` in a particular key path - - ``.find_all`` WILL NOT search the the value at the top-most level of the pytree (i.e. the root + - ``.find_all`` WILL NOT search the value at the top-most level of the pytree (i.e. the root node). The ``cond_fn`` will be evaluated recursively, starting at the root node's children. Example:: diff --git a/flax/linen/linear.py b/flax/linen/linear.py index 4f3592d54..27ace4869 100644 --- a/flax/linen/linear.py +++ b/flax/linen/linear.py @@ -76,7 +76,7 @@ class DenseGeneral(Module): >>> params = layer.init(jax.random.key(0), jnp.ones((1, 3))) >>> jax.tree_util.tree_map(jnp.shape, params) {'params': {'bias': (4, 5), 'kernel': (3, 4, 5)}} - >>> # apply transformation on the the second and last axes + >>> # apply transformation on the second and last axes >>> layer = nn.DenseGeneral(features=(4, 5), axis=(1, -1)) >>> params = layer.init(jax.random.key(0), jnp.ones((1, 3, 6, 7))) >>> jax.tree_util.tree_map(jnp.shape, params) diff --git a/flax/nnx/nn/linear.py b/flax/nnx/nn/linear.py index 10d516daf..d8e49023b 100644 --- a/flax/nnx/nn/linear.py +++ b/flax/nnx/nn/linear.py @@ -112,7 +112,7 @@ class LinearGeneral(Module): (2, 4, 5) >>> layer.bias.shape (4, 5) - >>> # apply transformation on the the second and last axes + >>> # apply transformation on the second and last axes >>> layer = nnx.LinearGeneral((2, 3), (4, 5), axis=(1, -1), rngs=nnx.Rngs(0)) >>> layer.kernel.shape (2, 3, 4, 5) diff --git a/flax/nnx/rnglib.py b/flax/nnx/rnglib.py index 833fbb9eb..1b87c0887 100644 --- a/flax/nnx/rnglib.py +++ b/flax/nnx/rnglib.py @@ -1525,7 +1525,7 @@ def reseed( NNX feature set including shared references. If ``False``, uses tree-mode which treats Modules as regular JAX pytrees, avoiding the overhead of the graph protocol. - policy: defines how the the new scalar key is for each RngStream is used to + policy: defines how the new scalar key is for each RngStream is used to reseed the stream. If ``'scalars_only'`` is given (the default), an error is raised if the target stream key is not a scalar. If ``'match_shape'`` is given, the new scalar key is split to match the shape of the target stream key. A callable diff --git a/flax/nnx/transforms/compilation.py b/flax/nnx/transforms/compilation.py index df69e1f2f..0f0e0787f 100644 --- a/flax/nnx/transforms/compilation.py +++ b/flax/nnx/transforms/compilation.py @@ -1385,7 +1385,7 @@ def f(m, x): jax.debug.visualize_array_sharding(y) Notice that here we simply used some ``PartitionSpec`` to define the spec - the the whole model and data. This works for simple cases but if we need + the whole model and data. This works for simple cases but if we need to assign different ``PartitionSpec`` to different parts of the model we need to use ``StateSharding`` and create some filters that allow us to target specific parts of the model. Here's an example of how to do tensor parallelism