Fix batch sizes in distributed_embedding_test.py.

[hertschuh]

Fixed inconsistencies between the per TPU batch size and global batch size.

[gemini-code-assist]

Summary of Changes

Hello @hertschuh, 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 addresses inconsistencies in batch size handling within distributed_embedding_test.py by centralizing the calculation of the global batch size. This ensures that tests for distributed embedding layers correctly account for the number of replicas in sync, improving the robustness and accuracy of the testing framework, especially for TPU environments.

Highlights

• Batch Size Standardization: The constant BATCH_SIZE_PER_CORE was renamed to BATCH_SIZE_PER_TPU for better semantic clarity, and a new self.batch_size instance variable was introduced to consistently represent the global batch size across all tests, calculated as BATCH_SIZE_PER_TPU * self._strategy.num_replicas_in_sync.
• JAX TPU Strategy Support: A JaxDummyStrategy class was added to correctly determine the number of replicas for the JAX backend when running on TPUs, ensuring proper batch size scaling in JAX environments.
• Input/Output Shape Alignment: All FeatureConfig definitions and tensor creations within the test suite now dynamically use self.batch_size for input and output shapes, eliminating hardcoded per-core batch sizes and ensuring consistency with the global batch size.
• Simplified Data Generation: The create_inputs_weights_and_labels helper function was refactored to remove the explicit batch_size parameter, instead deriving the batch size from the feature_config.input_shape, streamlining data generation for various input types (dense, ragged, sparse).

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 refactors the batch size handling in distributed_embedding_test.py to fix inconsistencies. The changes correctly use a global batch size throughout the tests, which is a good improvement. However, I've found a remaining inconsistency in the global batch size calculation between the TensorFlow and JAX backends within the test setup. My review includes suggestions to resolve this, ensuring the tests are robust and correct across different environments.

[abheesht17]

Thanks!