Adding get_device_count function to the distribution_lib #21791
Conversation
Summary of ChangesHello @buildwithsuhana, I'm Gemini Code Assist1! 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
Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
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 Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #21791 +/- ##
==========================================
+ Coverage 82.59% 82.63% +0.04%
==========================================
Files 572 577 +5
Lines 58535 59329 +794
Branches 9158 9300 +142
==========================================
+ Hits 48345 49026 +681
- Misses 7853 7915 +62
- Partials 2337 2388 +51
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
@buildwithsuhana |
| device_type: Optional device type to count (e.g., "cpu", | ||
| "gpu", "tpu"). If `None`, it counts all available | ||
| devices. |
There was a problem hiding this comment.
Oh, copy the one from the jax implementation. This one is incorrect.
5 participants
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.