Fix indices device error

[0x45f]

PR Category

Operator

Type of Change

Bug Fix

Description

Fix indices device error

Issue

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

[gemini-code-assist]

Summary of Changes

Hello @0x45f, 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 introduces a critical bug fix for the index_put and index_put_ functions within the flag_gems operator module. The primary purpose of this change is to prevent device-related errors that occur when the indices used in these operations are not located on the same computational device as the input tensor. By explicitly moving indices to the correct device, the robustness and reliability of these tensor manipulation functions are significantly improved, ensuring seamless execution across different hardware configurations.

Highlights

• Bug Fix: Resolved a device mismatch error in index_put and index_put_ operations by ensuring that all indices are moved to the same device as the input tensor (inp.device) before further processing.
• Device Consistency: Implemented a check and transfer mechanism for indices, preventing potential runtime errors when indices reside on a different device than the primary input tensor.

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 fixes a bug where indices tensors might be on a different device than the input tensor, causing an error. The fix correctly moves the indices tensors to the input tensor's device. My review focuses on improving the implementation by addressing code duplication and simplifying the logic. I've suggested refactoring the new code blocks in both index_put and index_put_ to improve maintainability and readability.

[gemini-code-assist]

This block can be simplified. The indices = list(indices) line is redundant because the following list comprehension creates a new list. You can remove the first line and the confusing double assignment to indices.

Additionally, this logic is duplicated in the index_put_ function. To improve maintainability, consider extracting this logic into a private helper function and calling it from both index_put and index_put_.

    indices = [
        index.to(inp.device) if index.device != inp.device else index
        for index in indices
    ]

[gemini-code-assist]

This block is identical to the one in index_put. To improve maintainability and avoid code duplication, consider extracting this logic into a private helper function.

As with the other occurrence, the indices = list(indices) line is redundant and the double assignment to indices is confusing. This can be simplified.

    indices = [
        index.to(inp.device) if index.device != inp.device else index
        for index in indices
    ]