Adding get_device_count function to the distribution_lib

[buildwithsuhana]

This PR introduces a new utility function, keras.distribution.get_device_count, to the Keras distribution API.
This function allows a user to query the total number of available devices (like GPUs, TPUs, or CPUs) that Keras can see.

[gemini-code-assist]

Summary of Changes

Hello @buildwithsuhana, 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 enhances Keras's device management capabilities by introducing a new "get_device_count" utility function within the "keras.distribution" module. This function provides a standardized and convenient way for users to programmatically determine the number of available computational devices, such as GPUs, TPUs, or CPUs, that Keras can leverage, thereby simplifying device configuration and resource allocation for distributed training.

Highlights

• New Utility Function: Introduced a new utility function, "keras.distribution.get_device_count", to the Keras distribution API.
• Device Query Capability: This function enables users to query the total number of available computational devices (GPUs, TPUs, or CPUs) visible to Keras.
• JAX Backend Implementation: Implemented the "get_device_count" functionality specifically for the JAX backend, utilizing "jax.device_count".
• API Exposure: Exposed the new function through the Keras public API in "keras.api._tf_keras.keras.distribution" and "keras.api.distribution".
• Test Case Update: Modified a JAX backend test to conditionally skip if the backend is not JAX or if fewer than 8 JAX devices are available, ensuring proper test environment.

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 introduces a new utility function get_device_count to the distribution API. The implementation looks good, but I've identified a couple of areas for improvement. The public API docstring for the new function is misleading regarding its behavior when no device type is specified, which could confuse users. I've suggested a clarification. Additionally, the new function lacks unit tests, which are essential for ensuring its correctness and preventing future regressions. I've recommended adding a test case.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 75.00000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.65%. Comparing base (4e1edcf) to head (0cfc5d7).
⚠️ Report is 37 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/api/_tf_keras/keras/distribution/__init__.py	0.00%	1 Missing ⚠️
keras/src/distribution/distribution_lib.py	66.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21791      +/-   ##
==========================================
+ Coverage   82.59%   82.65%   +0.06%     
==========================================
  Files         572      577       +5     
  Lines       58535    59335     +800     
  Branches     9158     9303     +145     
==========================================
+ Hits        48345    49044     +699     
- Misses       7853     7890      +37     
- Partials     2337     2401      +64     

Flag	Coverage Δ
keras	82.47% <75.00%> (+0.08%)	⬆️
keras-jax	63.34% <75.00%> (+0.15%)	⬆️
keras-numpy	57.57% <50.00%> (+0.01%)	⬆️
keras-openvino	34.29% <50.00%> (-0.05%)	⬇️
keras-tensorflow	64.13% <50.00%> (+0.18%)	⬆️
keras-torch	63.64% <50.00%> (+0.15%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[innat]

@buildwithsuhana
Thanks for the PR. It would be useful. There are some other utiliy (similar to this), would be great to have. #20744