From 88a51c304ec79016c9bcd2e58a2f5d06942599fd Mon Sep 17 00:00:00 2001 From: "Edward Z. Yang" Date: Wed, 12 Apr 2023 13:41:07 -0700 Subject: [PATCH 1/5] SymIntify randint and randperm Signed-off-by: Edward Z. Yang [ghstack-poisoned] --- aten/src/ATen/native/native_functions.yaml | 26 +++++++++++----------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index cb77cdba1821..7282478aa13a 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -4361,22 +4361,22 @@ CompositeExplicitAutograd: rand_like autogen: rand_like.out -- func: randint(int high, SymInt[] size, *, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor +- func: randint(SymInt high, SymInt[] size, *, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor tags: nondeterministic_seeded dispatch: CompositeExplicitAutograd: randint -- func: randint.generator(int high, SymInt[] size, *, Generator? generator, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor +- func: randint.generator(SymInt high, SymInt[] size, *, Generator? generator, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor tags: nondeterministic_seeded dispatch: CompositeExplicitAutograd: randint -- func: randint.low(int low, int high, SymInt[] size, *, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor +- func: randint.low(SymInt low, SymInt high, SymInt[] size, *, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor tags: nondeterministic_seeded dispatch: CompositeExplicitAutograd: randint -- func: randint.low_generator(int low, int high, SymInt[] size, *, Generator? generator, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor +- func: randint.low_generator(SymInt low, SymInt high, SymInt[] size, *, Generator? generator, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor tags: nondeterministic_seeded dispatch: CompositeExplicitAutograd: randint @@ -4386,22 +4386,22 @@ dispatch: CompositeExplicitAutograd: randint_out -- func: randint.generator_out(int high, SymInt[] size, *, Generator? generator, Tensor(a!) out) -> Tensor(a!) +- func: randint.generator_out(SymInt high, SymInt[] size, *, Generator? generator, Tensor(a!) out) -> Tensor(a!) tags: nondeterministic_seeded dispatch: CompositeExplicitAutograd: randint_out -- func: randint.low_out(int low, int high, SymInt[] size, *, Tensor(a!) out) -> Tensor(a!) +- func: randint.low_out(SymInt low, SymInt high, SymInt[] size, *, Tensor(a!) out) -> Tensor(a!) tags: nondeterministic_seeded dispatch: CompositeExplicitAutograd: randint_out -- func: randint.low_generator_out(int low, int high, SymInt[] size, *, Generator? generator, Tensor(a!) out) -> Tensor(a!) +- func: randint.low_generator_out(SymInt low, SymInt high, SymInt[] size, *, Generator? generator, Tensor(a!) out) -> Tensor(a!) tags: nondeterministic_seeded dispatch: CompositeExplicitAutograd: randint_out -- func: randint_like(Tensor self, int high, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor +- func: randint_like(Tensor self, SymInt high, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor tags: nondeterministic_seeded dispatch: # NB: Although this composite mutates on the inside, it is @@ -4409,7 +4409,7 @@ CompositeExplicitAutograd: randint_like autogen: randint_like.out -- func: randint_like.low_dtype(Tensor self, int low, int high, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor +- func: randint_like.low_dtype(Tensor self, SymInt low, SymInt high, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor tags: nondeterministic_seeded dispatch: # NB: Although this composite mutates on the inside, it is @@ -4457,22 +4457,22 @@ CompositeExplicitAutograd, CompositeImplicitAutogradNestedTensor: randn_like autogen: randn_like.out -- func: randperm(int n, *, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor +- func: randperm(SymInt n, *, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor tags: nondeterministic_seeded dispatch: CompositeExplicitAutograd: randperm -- func: randperm.generator(int n, *, Generator? generator, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor +- func: randperm.generator(SymInt n, *, Generator? generator, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor tags: nondeterministic_seeded dispatch: CompositeExplicitAutograd: randperm -- func: randperm.out(int n, *, Tensor(a!) out) -> Tensor(a!) +- func: randperm.out(SymInt n, *, Tensor(a!) out) -> Tensor(a!) tags: nondeterministic_seeded dispatch: CompositeExplicitAutograd: randperm_out -- func: randperm.generator_out(int n, *, Generator? generator, Tensor(a!) out) -> Tensor(a!) +- func: randperm.generator_out(SymInt n, *, Generator? generator, Tensor(a!) out) -> Tensor(a!) tags: nondeterministic_seeded dispatch: CPU: randperm_out_cpu From 9cf93a7f9ca401e946219f681e850a3be046418f Mon Sep 17 00:00:00 2001 From: "Edward Z. Yang" Date: Wed, 12 Apr 2023 13:43:23 -0700 Subject: [PATCH 2/5] Update on "SymIntify randint and randperm" Signed-off-by: Edward Z. Yang [ghstack-poisoned] --- aten/src/ATen/native/native_functions.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index 7282478aa13a..5186b45c0827 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -4381,7 +4381,7 @@ dispatch: CompositeExplicitAutograd: randint -- func: randint.out(int high, SymInt[] size, *, Tensor(a!) out) -> Tensor(a!) +- func: randint.out(SymInt high, SymInt[] size, *, Tensor(a!) out) -> Tensor(a!) tags: nondeterministic_seeded dispatch: CompositeExplicitAutograd: randint_out From 458071d65a6dad85a08e2776d70d06e2d2d517f0 Mon Sep 17 00:00:00 2001 From: "Edward Z. Yang" Date: Wed, 12 Apr 2023 13:47:43 -0700 Subject: [PATCH 3/5] Update on "SymIntify randint and randperm" Signed-off-by: Edward Z. Yang [ghstack-poisoned] --- aten/src/ATen/functorch/BatchRulesRandomness.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aten/src/ATen/functorch/BatchRulesRandomness.cpp b/aten/src/ATen/functorch/BatchRulesRandomness.cpp index c9482305bbd2..b1235b9f5689 100644 --- a/aten/src/ATen/functorch/BatchRulesRandomness.cpp +++ b/aten/src/ATen/functorch/BatchRulesRandomness.cpp @@ -298,10 +298,10 @@ struct RandIntBatchRuleHelper; template struct RandIntBatchRuleHelper> { - static Tensor apply(int64_t high, SymIntArrayRef shape, T... extra_args) { + static Tensor apply(c10::SymInt high, SymIntArrayRef shape, T... extra_args) { return random_batching_rule), &rand_int_wrapper, - int64_t, T...>(shape, high, std::forward(extra_args)...); + c10::SymInt, T...>(shape, high, std::forward(extra_args)...); } }; From d68c4a93bcb366153fc6c549fe5e6c00792d795f Mon Sep 17 00:00:00 2001 From: "Edward Z. Yang" Date: Wed, 12 Apr 2023 13:49:34 -0700 Subject: [PATCH 4/5] Update on "SymIntify randint and randperm" Signed-off-by: Edward Z. Yang [ghstack-poisoned] --- aten/src/ATen/functorch/BatchRulesRandomness.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aten/src/ATen/functorch/BatchRulesRandomness.cpp b/aten/src/ATen/functorch/BatchRulesRandomness.cpp index b1235b9f5689..7a67f516c587 100644 --- a/aten/src/ATen/functorch/BatchRulesRandomness.cpp +++ b/aten/src/ATen/functorch/BatchRulesRandomness.cpp @@ -279,7 +279,7 @@ struct RandomBatchRuleHelper> { }; template -Tensor rand_int_wrapper(SymIntArrayRef shape, int64_t high, T... extra_args) { +Tensor rand_int_wrapper(SymIntArrayRef shape, c10::SymInt high, T... extra_args) { return Func(high, shape, std::forward(extra_args)...); } From 9b0fe57a1cb151131b18cfb7b9ca95fe9c071f50 Mon Sep 17 00:00:00 2001 From: "Edward Z. Yang" Date: Wed, 12 Apr 2023 13:58:33 -0700 Subject: [PATCH 5/5] Update on "SymIntify randint and randperm" Signed-off-by: Edward Z. Yang [ghstack-poisoned] --- aten/src/ATen/functorch/BatchRulesRandomness.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aten/src/ATen/functorch/BatchRulesRandomness.cpp b/aten/src/ATen/functorch/BatchRulesRandomness.cpp index 7a67f516c587..f71931316992 100644 --- a/aten/src/ATen/functorch/BatchRulesRandomness.cpp +++ b/aten/src/ATen/functorch/BatchRulesRandomness.cpp @@ -318,7 +318,7 @@ struct RandTwoLeadingScalarsBatchRuleHelper> static Tensor apply(T0 scalar0, T1 scalar1, SymIntArrayRef shape, T... extra_args) { return random_batching_rule), &rand_int_low_wrapper, - int64_t, int64_t, T...>(shape, scalar0, scalar1, std::forward(extra_args)...); + T0, T1, T...>(shape, scalar0, scalar1, std::forward(extra_args)...); } };