Skip to content

Commit

Permalink
[FRONTEND] interleave does not need to check shape (#4535)
Browse files Browse the repository at this point in the history
1. Fixed the errors in the formulas in the comments.
2. When assert is enabled in DEBUG mode, using the list type will report
an error. Additionally, the Tensor returned by the internal interface
can ensure that the shape is a list and does not require checking.


The core Triton is a small number of people, and we receive many PRs
(thank
you!).  To help us review your code more quickly, **if you are a new
contributor (less than 3 PRs merged) we ask that you complete the
following
tasks and include the filled-out checklist in your PR description.**

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.

- [x] I am not making a trivial change, such as fixing a typo in a
comment.

- [x] I have written a PR description following these
  [rules](https://cbea.ms/git-commit/#why-not-how).

- [x] I have run `pre-commit run --from-ref origin/main --to-ref HEAD`.

- Select one of the following.
  - [ ] I have added tests.
    - `/test` for `lit` tests
    - `/unittest` for C++ tests
    - `/python/test` for end-to-end tests
- [x] This PR does not need a test because it does not take any effect
to working code.

- Select one of the following.
  - [x] I have not added any `lit` tests.
- [ ] The `lit` tests I have added follow these [best
practices](https://mlir.llvm.org/getting_started/TestingGuide/#filecheck-best-practices),
including the "tests should be minimal" section. (Usually running Python
code
    and using the instructions it generates is not minimal.)

Co-authored-by: 谢双镱 <xieshuangyi@bytedance.com>
  • Loading branch information
Mwsxy and 谢双镱 authored Aug 23, 2024
1 parent f210090 commit 54801fa
Showing 1 changed file with 1 addition and 2 deletions.
3 changes: 1 addition & 2 deletions python/triton/language/standard.py
Original file line number Diff line number Diff line change
Expand Up @@ -433,7 +433,7 @@ def flip(x, dim=None):
def interleave(a, b):
"""
Interleaves the values of two tensors along their last dimension. The two tensors must have the same shape.
Equivalent to `tl.join(a, b).reshape(a.shape[-1:] + [2 * a.shape[-1]])`
Equivalent to `tl.join(a, b).reshape(a.shape[:-1] + [2 * a.shape[-1]])`
:param a: The first input tensor.
:type a: Tensor
Expand All @@ -442,7 +442,6 @@ def interleave(a, b):
"""
c = core.join(a, b)

assert isinstance(c.shape, list)
if len(c.shape) == 1:
# We must have interleaved two scalars.
return c
Expand Down

0 comments on commit 54801fa

Please sign in to comment.