Skip to content

implement Repeat with fixed output shape#7114

Merged
alanwaketan merged 2 commits intomasterfrom
JackCaoG/repeat_static
May 25, 2024
Merged

implement Repeat with fixed output shape#7114
alanwaketan merged 2 commits intomasterfrom
JackCaoG/repeat_static

Conversation

@JackCaoG
Copy link
Copy Markdown
Collaborator

No description provided.

@JackCaoG JackCaoG changed the title torch_xla/experimental/custom_kernel.py implement Repeat with fixed output shape May 24, 2024
@JackCaoG JackCaoG marked this pull request as ready for review May 24, 2024 23:44
@JackCaoG JackCaoG requested a review from alanwaketan May 24, 2024 23:44

# shift the repeats by one
# tensor([0, 0, 1, 2, 0, 4, 0, 6, 7, 8])
exclusive_repeats = torch.roll(repeats, shifts=1)
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we want to shift right? If we don't shift, I assume we don't need to -1 at L700, and it will still output the correct answer?


# value in gather_indices represents the index in the input.
# tensor([1, 2, 2, 4, 4, 4, 4, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7])
gather_indices = torch.cumsum(block_split_indicators, dim=0) - 1
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Trying to reason about this step... but really hard to understand... Can you explain a little bit to me? I couldn't get why the block_split_indicators can be converted to gather_indices...

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can understand for each indicator, the index needs to bump. Trying to understand the value to bump.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay, now I understand. If some index are skipped, it means the indicator get selected more. The more it skips, the more it get selects. Which equals to the value to bump.

@alanwaketan
Copy link
Copy Markdown
Collaborator

Let's skip the GPU tests to move fast.

@alanwaketan alanwaketan merged commit 3369bf7 into master May 25, 2024
@alanwaketan alanwaketan deleted the JackCaoG/repeat_static branch May 25, 2024 00:36
@alanwaketan
Copy link
Copy Markdown
Collaborator

Thanks, Jack!

@alanwaketan alanwaketan requested a review from miladm May 25, 2024 01:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants