Commitf73345c

authored and

committed

Fix unsafe_chunk decomposition for empty tensors (#170053)

Fixes#147277Summary:- Fix `torch.unsafe_chunk` with empty size input tensor.- Matches eager mode behaviorPull Requestresolved:#170053Approved by:https://github.com/cyyever,https://github.com/jansel

1 parent9ec52b5 commitf73345cCopy full SHA for f73345c

File tree

+18

-1

lines changed

+18

-1

lines changed

Whitespace-only changes.

Lines changed: 14 additions & 0 deletions

Original file line number	Diff line number	Diff line change
`@@ -4592,6 +4592,20 @@ def fn2(a):`
`4592`	`4592`	`(torch.randn([2, 2, 10]),),`
`4593`	`4593`	`)`
`4594`	`4594`
	`4595`	`+ def test_unsafe_chunk_empty_tensor(self):`
	`4596`	`+ def fn(x, chunks):`
	`4597`	`+ return torch.unsafe_chunk(x, chunks=chunks, dim=0)`
	`4598`	`+`
	`4599`	`+ x = torch.zeros(0)`
	`4600`	`+ result = torch.compile(fn)(x, 4)`
	`4601`	`+ self.assertEqual(len(result), 4)`
	`4602`	`+ for chunk in result:`
	`4603`	`+ self.assertEqual(chunk.shape, torch.Size([0]))`
	`4604`	`+`
	`4605`	`+ self.common(fn, (torch.zeros(0), 4))`
	`4606`	`+ self.common(fn, (torch.zeros(0), 2))`
	`4607`	`+ self.common(fn, (torch.zeros(0, dtype=torch.float32), 3))`
	`4608`	`+`
`4595`	`4609`	`@parametrize("dilation", (1, 2))`
`4596`	`4610`	`@parametrize("dim", (subtest(2), subtest(3)))`
`4597`	`4611`	`def test_low_memory_max_pool(self, dilation: int, dim: int):`

Lines changed: 3 additions & 0 deletions

Original file line number	Diff line number	Diff line change
`@@ -187,6 +187,9 @@ def run(ex, *kwargs):`
`187`	`187`	`"test_empty1_dynamic_shapes":TestFailure(("cpu","cuda","xpu")),`
`188`	`188`	`"test_empty2_dynamic_shapes":TestFailure(("cpu","cuda","xpu")),`
`189`	`189`	`"test_empty_strided_dynamic_shapes":TestFailure(("cpu","cuda","xpu")),`
	`190`	`+"test_unsafe_chunk_empty_tensor_dynamic_shapes":TestFailure(`
	`191`	`+ ("cpu","cuda","xpu"),is_skip=True`
	`192`	`+ ),`
`190`	`193`	`"test_bucketize_nd_tiling_False_dynamic_shapes":TestFailure(("cpu",)),`
`191`	`194`	`"test_bucketize_nd_tiling_True_dynamic_shapes":TestFailure(("cpu",)),`
`192`	`195`	`"test_bucketize_default_kwargs_dynamic_shapes":TestFailure(("cpu",)),`

Lines changed: 1 addition & 1 deletion

Original file line number	Diff line number	Diff line change
`@@ -2040,7 +2040,7 @@ def unsafe_chunk_py_impl(tensor, chunks, dim=0) -> list[Tensor]:`
`2040`	`2040`	`split_size= (dim_size+chunks-1)//chunks`
`2041`	`2041`
`2042`	`2042`	`ifsplit_size==0anddim_size==0:`
`2043`		`-split_sizes= [split_sizefor_inchunks]`
	`2043`	`+split_sizes= [split_sizefor_inrange(chunks)]`
`2044`	`2044`	`split_sizes[chunks-1]=split_size- (split_size*chunks-dim_size)`
`2045`	`2045`	`returntorch.ops.aten.unsafe_split_with_sizes.default(tensor,split_sizes,dim)`
`2046`	`2046`	`returntorch.ops.aten.unsafe_split.Tensor(tensor,split_size,dim)`

Comments

(0)