Skip to content

[Text Pipeline] Implement Feature Extractor#337

Merged
copybara-service[bot] merged 1 commit intoAI-Hypercomputer:mainfrom
syhuang22:feat/ltx2-feature-extractor
Mar 2, 2026
Merged

[Text Pipeline] Implement Feature Extractor#337
copybara-service[bot] merged 1 commit intoAI-Hypercomputer:mainfrom
syhuang22:feat/ltx2-feature-extractor

Conversation

@syhuang22
Copy link
Collaborator

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.

@syhuang22 syhuang22 requested a review from entrpn as a code owner February 26, 2026 00:22
@google-cla
Copy link

google-cla bot commented Feb 26, 2026

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.

@syhuang22 syhuang22 force-pushed the feat/ltx2-feature-extractor branch 4 times, most recently from 5ecb3fe to 3774a5f Compare February 26, 2026 01:24
entrpn
entrpn previously requested changes Feb 26, 2026
@entrpn entrpn dismissed their stale review February 27, 2026 20:05

make clear where the value is coming from.

entrpn
entrpn previously approved these changes Feb 27, 2026
@entrpn
Copy link
Collaborator

entrpn commented Feb 27, 2026

@syhuang22 please squash your commits.

@syhuang22 syhuang22 force-pushed the feat/ltx2-feature-extractor branch 4 times, most recently from 27cf2f9 to ee6d38d Compare March 2, 2026 18:48
@syhuang22 syhuang22 force-pushed the feat/ltx2-feature-extractor branch from ee6d38d to ce7d1f4 Compare March 2, 2026 20:55
Signed-off-by: James Huang <syhuang1201@gmail.com>
@syhuang22 syhuang22 force-pushed the feat/ltx2-feature-extractor branch from ce7d1f4 to f9f3e89 Compare March 2, 2026 21:00
@copybara-service copybara-service bot merged commit 20d650a into AI-Hypercomputer:main Mar 2, 2026
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants