Skip to content

This CL introduces an end-to-end Flax implementation of the DLRM-HSTU model, along with individual unit tests for the various components.#64

Merged
copybara-service[bot] merged 1 commit into
mainfrom
test_789375174
Jun 11, 2026
Merged

This CL introduces an end-to-end Flax implementation of the DLRM-HSTU model, along with individual unit tests for the various components.#64
copybara-service[bot] merged 1 commit into
mainfrom
test_789375174

Conversation

@copybara-service

@copybara-service copybara-service Bot commented Jul 31, 2025

Copy link
Copy Markdown

This CL introduces an end-to-end Flax implementation of the DLRM-HSTU model, along with individual unit tests for the various components.

  • Converted all necessary PyTorch modules to Flax/Linen
  • Full training loop test for end to end model
  • Added unit tests for individual components (e.g. ActionEncoder, ContentEncoder, STU modules)
  • Ensured model runs correctly on TPUs.
  • Verified shape-correctness of all modules and parameters.

Reverts changelist 793734230

@copybara-service copybara-service Bot force-pushed the test_789375174 branch 5 times, most recently from 19b260b to e1b9176 Compare August 22, 2025 00:05
Reverts changelist 793734230

PiperOrigin-RevId: 930775912
@copybara-service copybara-service Bot merged commit 8489b48 into main Jun 11, 2026
89 of 115 checks passed
@copybara-service copybara-service Bot deleted the test_789375174 branch June 11, 2026 22:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant