Fix doubled 'the the' typos across multiple files

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::

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)

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)

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

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