def __call__(
        self,
        x,
        hidden_forward=None,
        cell_forward=None,
        hidden_backward=None,
        cell_backward=None,
    ):
        """
        Process input sequence in both directions and concatenate the results.

        Args:
            x: Input tensor of shape (batch_size, seq_len, input_size)
            hidden_forward: Initial hidden state for forward direction
            cell_forward: Initial cell state for forward direction
            hidden_backward: Initial hidden state for backward direction
            cell_backward: Initial cell state for backward direction

        Returns:
            Tuple of:
                - Combined output hidden states (batch_size, seq_len, 2*hidden_size)
                - Tuple of final states ((forward_hidden, forward_cell), (backward_hidden, backward_cell))
        """

        if x.ndim == 2:
            x = mx.expand_dims(x, axis=0)  # (1, seq_len, input_size)

        # Forward direction
        forward_hidden, forward_cell = self._forward_direction(
            x, hidden_forward, cell_forward
        )

        # Backward direction
        backward_hidden, backward_cell = self._backward_direction(
            x, hidden_backward, cell_backward
        )

        # Concatenate outputs along the feature dimension
        output = mx.concatenate([forward_hidden, backward_hidden], axis=-1)

        # Return combined output and final states for both directions
        return output, (
            (forward_hidden[..., -1, :], forward_cell[..., -1, :]),
            (backward_hidden[..., 0, :], backward_cell[..., 0, :]),
        )

        Process input sequence in both directions and concatenate the results.

        Args:
            x: Input tensor of shape (batch_size, seq_len, input_size)
            hidden_forward: Initial hidden state for forward direction
            cell_forward: Initial cell state for forward direction
            hidden_backward: Initial hidden state for backward direction
            cell_backward: Initial cell state for backward direction

        Returns:
            Tuple of:
                - Combined output hidden states (batch_size, seq_len, 2*hidden_size)
                - Tuple of final states ((forward_hidden, forward_cell), (backward_hidden, backward_cell))
        
mlx_audio/tts/models/kokoro/modules.py