Expand Embed test coverage with functional, attend, and dtype tests

[us]

What does this PR do?

Expands test coverage for nnx.Embed (tests/nnx/nn/embed_test.py), which previously had only 77 lines with a single Linen consistency test class.

Changes

New TestEmbed class — Functional tests:

• Output shape verification (1D and 2D inputs)
• Embedding parameter type (nnx.Param) and shape
• ValueError on float input
• num_embeddings=1 broadcast behavior
• Custom embedding_init (zeros)
• Negative indexing
• Parameterized dtype/param_dtype combinations
• dtype=None defaults to param_dtype

New TestEmbedAttend class — attend() method tests:

• Output shape verification
• Parameterized computation correctness across dtype combinations

Extended TestLinenConsistency class:

• Fixed existing test: replaced assert_array_equal with assert_allclose + conditional rtol for float16
• Added _make_models helper to reduce duplication
• Added guard assertions for parameter equivalence
• New: attend() Linen/NNX consistency (parameterized)
• New: jax.value_and_grad gradient consistency (parameterized)

Test count: 16 → 43

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other checks if that's the case).
• This change is discussed in a Github issue/discussion (please add a link).
• The documentation and docstrings adhere to the documentation guidelines.
• This change includes necessary high-coverage tests. (No quality testing = no merge!)

[gemini-code-assist]

Summary of Changes

Hello @us, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request substantially enhances the test coverage for the nnx.Embed module. It introduces new functional tests to validate core behaviors like output shapes, parameter properties, and error conditions, alongside dedicated tests for the attend() method. Furthermore, it strengthens the existing Linen consistency tests by adding attend() and gradient consistency checks, ensuring nnx.Embed behaves as expected across various data types and aligns with its Flax Linen counterpart.

Highlights

• Expanded Functional Tests for nnx.Embed: Introduced a new TestEmbed class covering output shapes (1D/2D), parameter type and shape, error handling for invalid input dtypes, broadcast behavior for single embeddings, custom initializers, negative indexing, and comprehensive dtype combinations.
• Added attend() Method Tests: A new TestEmbedAttend class was created to specifically test the attend() method, verifying its output shape and computational correctness across various dtype configurations.
• Enhanced Linen Consistency and Gradient Tests: The TestLinenConsistency class was refactored with a _make_models helper and updated to use assert_allclose for better float comparison. New tests were added for attend() method consistency between NNX and Linen, and for jax.value_and_grad gradient consistency, ensuring robust comparisons across different dtypes.
• Increased Test Coverage: The total number of tests for nnx.Embed significantly increased from 16 to 43, improving the overall reliability of the component.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• tests/nnx/nn/embed_test.py
  • Added TestEmbed class with functional tests for nnx.Embed.
  • Added TestEmbedAttend class for testing the attend() method.
  • Refactored TestLinenConsistency with a _make_models helper.
  • Updated TestLinenConsistency to use assert_allclose and added new tests for attend() and gradient consistency.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request significantly expands the test coverage for nnx.Embed, which is a valuable improvement. It introduces new functional tests, tests for the attend method, and enhances the Linen consistency tests, including crucial gradient checks. The refactoring of existing tests to use assert_allclose for float comparisons and a helper method for model creation improves both correctness and maintainability. I have one suggestion to further reduce code duplication.