Document LSTMCell carry as the (c, h) tuple it actually is

The LSTMCell.__call__ docstring described 'carry' as 'the hidden state
of the LSTM cell', leaving users to infer that it is actually a
tuple of cell state and hidden state both of shape (*batch, features),
typically created via LSTMCell.initialize_carry. Spell that contract
out in both the Linen and NNX twins, and mirror the same docstring
for OptimizedLSTMCell.

Fixes #4124

Author: adityasingh2400

flax/linen/recurrent.py

@@ -136,7 +136,8 @@ def __call__(self, carry, inputs):
     r"""A long short-term memory (LSTM) cell.
 
     Args:
-      carry: the hidden state of the LSTM cell,
+      carry: a tuple ``(c, h)`` of the cell state ``c`` and the hidden
+        state ``h``, both of shape ``(*batch, features)``. Typically
         initialized using ``LSTMCell.initialize_carry``.
       inputs: an ndarray with the input for the current time step.
         All dimensions except the final are considered batch dimensions.
@@ -285,8 +286,9 @@ def __call__(
     r"""An optimized long short-term memory (LSTM) cell.
 
     Args:
-      carry: the hidden state of the LSTM cell, initialized using
-        ``LSTMCell.initialize_carry``.
+      carry: a tuple ``(c, h)`` of the cell state ``c`` and the hidden
+        state ``h``, both of shape ``(*batch, features)``. Typically
+        initialized using ``OptimizedLSTMCell.initialize_carry``.
       inputs: an ndarray with the input for the current time step. All
         dimensions except the final are considered batch dimensions.
 

flax/nnx/nn/recurrent.py

@@ -201,7 +201,8 @@ def __call__(
     r"""A long short-term memory (LSTM) cell.
 
     Args:
-      carry: the hidden state of the LSTM cell,
+      carry: a tuple ``(c, h)`` of the cell state ``c`` and the hidden
+        state ``h``, both of shape ``(*batch, features)``. Typically
         initialized using ``LSTMCell.initialize_carry``.
       inputs: an ndarray with the input for the current time step.
         All dimensions except the final are considered batch dimensions.
@@ -382,8 +383,9 @@ def __call__(
     r"""An optimized long short-term memory (LSTM) cell.
 
     Args:
-      carry: the hidden state of the LSTM cell, initialized using
-        ``LSTMCell.initialize_carry``.
+      carry: a tuple ``(c, h)`` of the cell state ``c`` and the hidden
+        state ``h``, both of shape ``(*batch, features)``. Typically
+        initialized using ``OptimizedLSTMCell.initialize_carry``.
       inputs: an ndarray with the input for the current time step.
         All dimensions except the final are considered batch dimensions.