[Text Pipeline] Implement Feature Extractor

[syhuang22]

This is the feature extractor of LTX2, which maps the flattened hidden states from the Gemma 3 text encoder to the embedding dimension.

The extraction process has been implemented in JAX/Flax following the exact three sequential steps required by the LTX-2 architecture:

1. Mean-Centered Scaling: Applied to the intermediate outputs across both the sequence and embedding dimensions for each layer to normalize the features.
2. Flattening: The scaled outputs from all layers (L) are stacked and flattened into a high-dimensional representation of shape [B, T, D × L] (e.g., [B, T, 3840 × 49]).
3. Linear Projection: Added a JAX-equivalent linear projection layer (torch.nn.Linear(3840 * 49, 3840)) to map the high-dimensional representation down to the target embedding dimension D.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

make clear where the value is coming from.

[entrpn]

@syhuang22 please squash your commits.