From d314dfed11e342791b83bbb0cb38763bf7265717 Mon Sep 17 00:00:00 2001 From: Abhinav-Khot <23110006@iitgn.ac.in> Date: Tue, 18 Mar 2025 23:24:28 +0530 Subject: [PATCH 1/6] Modify arange to directly determine the shape from the arguments if possible --- pytensor/tensor/basic.py | 18 +++++++++++++++++- tests/tensor/test_basic.py | 6 ++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index e30887cfe3..12ef2e1b77 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -3224,13 +3224,29 @@ def __init__(self, dtype): self.dtype = dtype def make_node(self, start, stop, step): + types = TensorConstant + shape = (None,) + # if it is possible to directly determine the shape i.e static shape is present, we find it. + if ( + isinstance(start, types) + and isinstance(stop, types) + and isinstance(step, types) + ): + length = max( + np.max(np.ceil((stop.value - start.value) / step.value)) + .astype(int) + .item(), + 0, + ) + shape = (length,) + start, stop, step = map(as_tensor_variable, (start, stop, step)) assert start.ndim == 0 assert stop.ndim == 0 assert step.ndim == 0 inputs = [start, stop, step] - outputs = [tensor(dtype=self.dtype, shape=(None,))] + outputs = [tensor(dtype=self.dtype, shape=shape)] return Apply(self, inputs, outputs) diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py index 60643e2984..29ad4c9458 100644 --- a/tests/tensor/test_basic.py +++ b/tests/tensor/test_basic.py @@ -2861,6 +2861,12 @@ def test_infer_shape(self, cast_policy): assert np.all(f(2) == len(np.arange(0, 2))) assert np.all(f(0) == len(np.arange(0, 0))) + def test_static_shape(self): + assert np.arange(1, 10).shape == arange(1, 10).type.shape + assert np.arange(10, 1, -1).shape == arange(10, 1, -1).type.shape + assert np.arange(1, -9, 2).shape == arange(1, -9, 2).type.shape + assert np.arange(1.3, 17.48, 2.67).shape == arange(1.3, 17.48, 2.67).type.shape + class TestNdGrid: def setup_method(self): From bb5601073fdadd600b4d3760b931f83001d8f0cd Mon Sep 17 00:00:00 2001 From: Abhinav-Khot <23110006@iitgn.ac.in> Date: Wed, 19 Mar 2025 00:19:53 +0530 Subject: [PATCH 2/6] Fix overflow issue --- pytensor/tensor/basic.py | 8 ++++---- tests/tensor/test_basic.py | 1 + 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index 12ef2e1b77..76e87f7808 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -3224,6 +3224,8 @@ def __init__(self, dtype): self.dtype = dtype def make_node(self, start, stop, step): + from math import ceil + types = TensorConstant shape = (None,) # if it is possible to directly determine the shape i.e static shape is present, we find it. @@ -3232,11 +3234,9 @@ def make_node(self, start, stop, step): and isinstance(stop, types) and isinstance(step, types) ): + # Convert to int32 or int64 before calculation length = max( - np.max(np.ceil((stop.value - start.value) / step.value)) - .astype(int) - .item(), - 0, + ceil((float(stop.value) - float(start.value)) / float(step.value)), 0 ) shape = (length,) diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py index 29ad4c9458..19e08d49cd 100644 --- a/tests/tensor/test_basic.py +++ b/tests/tensor/test_basic.py @@ -2866,6 +2866,7 @@ def test_static_shape(self): assert np.arange(10, 1, -1).shape == arange(10, 1, -1).type.shape assert np.arange(1, -9, 2).shape == arange(1, -9, 2).type.shape assert np.arange(1.3, 17.48, 2.67).shape == arange(1.3, 17.48, 2.67).type.shape + assert np.arange(-64, 64).shape == arange(-64, 64).type.shape class TestNdGrid: From 65078a5e07606c7e523f051f027740bdc6a4ae95 Mon Sep 17 00:00:00 2001 From: Abhinav-Khot <23110006@iitgn.ac.in> Date: Wed, 19 Mar 2025 10:31:19 +0530 Subject: [PATCH 3/6] removed redundant variable --- pytensor/tensor/basic.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index 76e87f7808..5223b05f46 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -3226,17 +3226,15 @@ def __init__(self, dtype): def make_node(self, start, stop, step): from math import ceil - types = TensorConstant shape = (None,) # if it is possible to directly determine the shape i.e static shape is present, we find it. if ( - isinstance(start, types) - and isinstance(stop, types) - and isinstance(step, types) + isinstance(start, TensorConstant) + and isinstance(stop, TensorConstant) + and isinstance(step, TensorConstant) ): - # Convert to int32 or int64 before calculation length = max( - ceil((float(stop.value) - float(start.value)) / float(step.value)), 0 + ceil((float(stop) - float(start.value)) / float(step.value)), 0 ) shape = (length,) From 098541fc5212012d02221bf39167d7cb3a579735 Mon Sep 17 00:00:00 2001 From: Abhinav-Khot <23110006@iitgn.ac.in> Date: Wed, 19 Mar 2025 11:57:50 +0530 Subject: [PATCH 4/6] Fix minor issue --- pytensor/tensor/basic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index 5223b05f46..414f9a970a 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -3234,7 +3234,7 @@ def make_node(self, start, stop, step): and isinstance(step, TensorConstant) ): length = max( - ceil((float(stop) - float(start.value)) / float(step.value)), 0 + ceil((float(stop.value) - float(start.value)) / float(step.value)), 0 ) shape = (length,) From ebb0e75753b6ee8aa5bc502fe09da3367e1062a3 Mon Sep 17 00:00:00 2001 From: Abhinav-Khot <23110006@iitgn.ac.in> Date: Wed, 19 Mar 2025 16:51:40 +0530 Subject: [PATCH 5/6] Modify make_node for completeness and clarity --- pytensor/tensor/basic.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index 414f9a970a..bbed78dc4e 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -3226,6 +3226,12 @@ def __init__(self, dtype): def make_node(self, start, stop, step): from math import ceil + start, stop, step = map(as_tensor_variable, (start, stop, step)) + + assert start.ndim == 0 + assert stop.ndim == 0 + assert step.ndim == 0 + shape = (None,) # if it is possible to directly determine the shape i.e static shape is present, we find it. if ( @@ -3234,15 +3240,10 @@ def make_node(self, start, stop, step): and isinstance(step, TensorConstant) ): length = max( - ceil((float(stop.value) - float(start.value)) / float(step.value)), 0 + ceil((float(stop.data) - float(start.data)) / float(step.data)), 0 ) shape = (length,) - start, stop, step = map(as_tensor_variable, (start, stop, step)) - assert start.ndim == 0 - assert stop.ndim == 0 - assert step.ndim == 0 - inputs = [start, stop, step] outputs = [tensor(dtype=self.dtype, shape=shape)] From 11a915bcbea7a3901016b00c72344d3108bd0cd5 Mon Sep 17 00:00:00 2001 From: Abhinav-Khot <23110006@iitgn.ac.in> Date: Thu, 20 Mar 2025 18:06:51 +0530 Subject: [PATCH 6/6] minor fix for readability --- pytensor/tensor/basic.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index bbed78dc4e..edd37b4af5 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -3232,7 +3232,6 @@ def make_node(self, start, stop, step): assert stop.ndim == 0 assert step.ndim == 0 - shape = (None,) # if it is possible to directly determine the shape i.e static shape is present, we find it. if ( isinstance(start, TensorConstant) @@ -3243,6 +3242,8 @@ def make_node(self, start, stop, step): ceil((float(stop.data) - float(start.data)) / float(step.data)), 0 ) shape = (length,) + else: + shape = (None,) inputs = [start, stop, step] outputs = [tensor(dtype=self.dtype, shape=shape)]