From b51e04cc1ef358979283a3816c2c008face76c59 Mon Sep 17 00:00:00 2001 From: Ben Arthur Date: Thu, 16 Apr 2026 10:38:01 -0400 Subject: [PATCH 1/9] add size(A,d) methods --- src/TensorStoreWrapper.jl | 2 ++ test/runtests.jl | 6 ++++++ 2 files changed, 8 insertions(+) diff --git a/src/TensorStoreWrapper.jl b/src/TensorStoreWrapper.jl index 160fc34..cd5f6b1 100644 --- a/src/TensorStoreWrapper.jl +++ b/src/TensorStoreWrapper.jl @@ -150,6 +150,7 @@ function Base.setindex!(w::TensorStoreWrapper, v, indices...; kwargs...) end Base.size(w::TensorStoreWrapper) = pyconvert(Tuple, parent(w).shape) +Base.size(w::TensorStoreWrapper, d::Integer) = d <= ndims(w) ? size(w)[d] : 1 Base.ndims(w::TensorStoreWrapper) = pyconvert(Int, parent(w).rank) const TS_TYPE_MAP = Dict( @@ -228,6 +229,7 @@ end # IndexDomainWrapper methods Base.size(w::IndexDomainWrapper) = pyconvert(Tuple, parent(w).shape) +Base.size(w::IndexDomainWrapper, d::Integer) = d <= ndims(w) ? size(w)[d] : 1 Base.ndims(w::IndexDomainWrapper) = pyconvert(Int, parent(w).rank) function Base.axes(w::IndexDomainWrapper) rank = ndims(w) diff --git a/test/runtests.jl b/test/runtests.jl index 4068501..2684c5b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -23,6 +23,9 @@ using PythonCall @test eltype(w) == Int32 @test ndims(w) == 2 @test size(w) == (10, 20) + @test size(w, 1) == 10 + @test size(w, 2) == 20 + @test size(w, 3) == 1 @test axes(w) == (1:10, 1:20) end @@ -44,6 +47,9 @@ using PythonCall # Labeled indexing sub_w = w[x=1:5, y=11:15] @test size(sub_w) == (5, 5) + @test size(domain, 1) == 10 + @test size(domain, 2) == 20 + @test size(domain, 3) == 1 @test axes(sub_w) == (1:5, 11:15) # translate_by From b8c8183de38573ca74e3fd8a10b0f6bf058ea382 Mon Sep 17 00:00:00 2001 From: Ben Arthur Date: Thu, 16 Apr 2026 12:34:00 -0400 Subject: [PATCH 2/9] handle d<1 Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/TensorStoreWrapper.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/TensorStoreWrapper.jl b/src/TensorStoreWrapper.jl index cd5f6b1..2138045 100644 --- a/src/TensorStoreWrapper.jl +++ b/src/TensorStoreWrapper.jl @@ -150,7 +150,11 @@ function Base.setindex!(w::TensorStoreWrapper, v, indices...; kwargs...) end Base.size(w::TensorStoreWrapper) = pyconvert(Tuple, parent(w).shape) -Base.size(w::TensorStoreWrapper, d::Integer) = d <= ndims(w) ? size(w)[d] : 1 +function Base.size(w::TensorStoreWrapper, d::Integer) + d < 1 && throw(ArgumentError("dimension must be ≥ 1")) + sz = size(w) + return d <= length(sz) ? sz[d] : 1 +end Base.ndims(w::TensorStoreWrapper) = pyconvert(Int, parent(w).rank) const TS_TYPE_MAP = Dict( From 67666bf06782efb9edfb67d2d5aabaae31000d9d Mon Sep 17 00:00:00 2001 From: Ben Arthur Date: Thu, 16 Apr 2026 12:34:13 -0400 Subject: [PATCH 3/9] handle d<1 Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/TensorStoreWrapper.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/TensorStoreWrapper.jl b/src/TensorStoreWrapper.jl index 2138045..6d7e67e 100644 --- a/src/TensorStoreWrapper.jl +++ b/src/TensorStoreWrapper.jl @@ -233,7 +233,11 @@ end # IndexDomainWrapper methods Base.size(w::IndexDomainWrapper) = pyconvert(Tuple, parent(w).shape) -Base.size(w::IndexDomainWrapper, d::Integer) = d <= ndims(w) ? size(w)[d] : 1 +function Base.size(w::IndexDomainWrapper, d::Integer) + d < 1 && throw(ArgumentError("dimension must be ≥ 1")) + sz = size(w) + return d <= length(sz) ? sz[d] : 1 +end Base.ndims(w::IndexDomainWrapper) = pyconvert(Int, parent(w).rank) function Base.axes(w::IndexDomainWrapper) rank = ndims(w) From d6dc0b581ec464e6f7ee3d6670c56ddc28d90426 Mon Sep 17 00:00:00 2001 From: Ben Arthur Date: Thu, 16 Apr 2026 12:39:41 -0400 Subject: [PATCH 4/9] correct check for ndims --- src/TensorStoreWrapper.jl | 6 ++---- test/runtests.jl | 2 ++ 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/TensorStoreWrapper.jl b/src/TensorStoreWrapper.jl index 6d7e67e..5541cc4 100644 --- a/src/TensorStoreWrapper.jl +++ b/src/TensorStoreWrapper.jl @@ -152,8 +152,7 @@ end Base.size(w::TensorStoreWrapper) = pyconvert(Tuple, parent(w).shape) function Base.size(w::TensorStoreWrapper, d::Integer) d < 1 && throw(ArgumentError("dimension must be ≥ 1")) - sz = size(w) - return d <= length(sz) ? sz[d] : 1 + return d <= ndims(w) ? size(w)[d] : 1 end Base.ndims(w::TensorStoreWrapper) = pyconvert(Int, parent(w).rank) @@ -235,8 +234,7 @@ end Base.size(w::IndexDomainWrapper) = pyconvert(Tuple, parent(w).shape) function Base.size(w::IndexDomainWrapper, d::Integer) d < 1 && throw(ArgumentError("dimension must be ≥ 1")) - sz = size(w) - return d <= length(sz) ? sz[d] : 1 + return d <= ndims(w) ? size(w)[d] : 1 end Base.ndims(w::IndexDomainWrapper) = pyconvert(Int, parent(w).rank) function Base.axes(w::IndexDomainWrapper) diff --git a/test/runtests.jl b/test/runtests.jl index 2684c5b..21727b4 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -26,6 +26,7 @@ using PythonCall @test size(w, 1) == 10 @test size(w, 2) == 20 @test size(w, 3) == 1 + @test_throws ArgumentError size(w, 0) @test axes(w) == (1:10, 1:20) end @@ -50,6 +51,7 @@ using PythonCall @test size(domain, 1) == 10 @test size(domain, 2) == 20 @test size(domain, 3) == 1 + @test_throws ArgumentError size(domain, 0) @test axes(sub_w) == (1:5, 11:15) # translate_by From 13b6f46db15b5a8b13c07508229d24fb3ff70893 Mon Sep 17 00:00:00 2001 From: Ben Arthur Date: Tue, 21 Apr 2026 17:24:49 -0400 Subject: [PATCH 5/9] Update src/TensorStoreWrapper.jl Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/TensorStoreWrapper.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/TensorStoreWrapper.jl b/src/TensorStoreWrapper.jl index 5541cc4..04c1510 100644 --- a/src/TensorStoreWrapper.jl +++ b/src/TensorStoreWrapper.jl @@ -152,7 +152,8 @@ end Base.size(w::TensorStoreWrapper) = pyconvert(Tuple, parent(w).shape) function Base.size(w::TensorStoreWrapper, d::Integer) d < 1 && throw(ArgumentError("dimension must be ≥ 1")) - return d <= ndims(w) ? size(w)[d] : 1 + di = Base.to_index(d) + return di <= ndims(w) ? size(w)[di] : 1 end Base.ndims(w::TensorStoreWrapper) = pyconvert(Int, parent(w).rank) From 86d00455dfcde9ab8200ad903b8968616d65e759 Mon Sep 17 00:00:00 2001 From: Ben Arthur Date: Tue, 21 Apr 2026 17:24:57 -0400 Subject: [PATCH 6/9] Update src/TensorStoreWrapper.jl Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/TensorStoreWrapper.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/TensorStoreWrapper.jl b/src/TensorStoreWrapper.jl index 04c1510..797fa21 100644 --- a/src/TensorStoreWrapper.jl +++ b/src/TensorStoreWrapper.jl @@ -235,7 +235,8 @@ end Base.size(w::IndexDomainWrapper) = pyconvert(Tuple, parent(w).shape) function Base.size(w::IndexDomainWrapper, d::Integer) d < 1 && throw(ArgumentError("dimension must be ≥ 1")) - return d <= ndims(w) ? size(w)[d] : 1 + di = Base.to_index(d) + return d <= ndims(w) ? size(w)[di] : 1 end Base.ndims(w::IndexDomainWrapper) = pyconvert(Int, parent(w).rank) function Base.axes(w::IndexDomainWrapper) From 6d9713831bff50d75da993ffe2c0af605845278d Mon Sep 17 00:00:00 2001 From: Ben Arthur Date: Tue, 21 Apr 2026 17:27:10 -0400 Subject: [PATCH 7/9] presumably less memory allocations --- src/TensorStoreWrapper.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/TensorStoreWrapper.jl b/src/TensorStoreWrapper.jl index 797fa21..f2619a3 100644 --- a/src/TensorStoreWrapper.jl +++ b/src/TensorStoreWrapper.jl @@ -153,7 +153,7 @@ Base.size(w::TensorStoreWrapper) = pyconvert(Tuple, parent(w).shape) function Base.size(w::TensorStoreWrapper, d::Integer) d < 1 && throw(ArgumentError("dimension must be ≥ 1")) di = Base.to_index(d) - return di <= ndims(w) ? size(w)[di] : 1 + return d <= ndims(w) ? pyconvert(Int, parent(w).shape[di-1]) : 1 end Base.ndims(w::TensorStoreWrapper) = pyconvert(Int, parent(w).rank) @@ -236,7 +236,7 @@ Base.size(w::IndexDomainWrapper) = pyconvert(Tuple, parent(w).shape) function Base.size(w::IndexDomainWrapper, d::Integer) d < 1 && throw(ArgumentError("dimension must be ≥ 1")) di = Base.to_index(d) - return d <= ndims(w) ? size(w)[di] : 1 + return d <= ndims(w) ? pyconvert(Int, parent(w).shape[di-1]) : 1 end Base.ndims(w::IndexDomainWrapper) = pyconvert(Int, parent(w).rank) function Base.axes(w::IndexDomainWrapper) From 6e58653f914d6dc1a7f7c724d47626461ee55569 Mon Sep 17 00:00:00 2001 From: Mark Kittisopikul Date: Thu, 23 Apr 2026 02:10:19 -0400 Subject: [PATCH 8/9] Update src/TensorStoreWrapper.jl Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/TensorStoreWrapper.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/TensorStoreWrapper.jl b/src/TensorStoreWrapper.jl index f2619a3..2a83624 100644 --- a/src/TensorStoreWrapper.jl +++ b/src/TensorStoreWrapper.jl @@ -152,8 +152,9 @@ end Base.size(w::TensorStoreWrapper) = pyconvert(Tuple, parent(w).shape) function Base.size(w::TensorStoreWrapper, d::Integer) d < 1 && throw(ArgumentError("dimension must be ≥ 1")) + d > ndims(w) && return 1 di = Base.to_index(d) - return d <= ndims(w) ? pyconvert(Int, parent(w).shape[di-1]) : 1 + return pyconvert(Int, parent(w).shape[di-1]) end Base.ndims(w::TensorStoreWrapper) = pyconvert(Int, parent(w).rank) From 911282de3b40464e39b8ebac3d4a2cf70de7dd24 Mon Sep 17 00:00:00 2001 From: Mark Kittisopikul Date: Thu, 23 Apr 2026 02:10:29 -0400 Subject: [PATCH 9/9] Update src/TensorStoreWrapper.jl Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/TensorStoreWrapper.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/TensorStoreWrapper.jl b/src/TensorStoreWrapper.jl index 2a83624..da888a4 100644 --- a/src/TensorStoreWrapper.jl +++ b/src/TensorStoreWrapper.jl @@ -236,8 +236,9 @@ end Base.size(w::IndexDomainWrapper) = pyconvert(Tuple, parent(w).shape) function Base.size(w::IndexDomainWrapper, d::Integer) d < 1 && throw(ArgumentError("dimension must be ≥ 1")) + d > ndims(w) && return 1 di = Base.to_index(d) - return d <= ndims(w) ? pyconvert(Int, parent(w).shape[di-1]) : 1 + return pyconvert(Int, parent(w).shape[di-1]) end Base.ndims(w::IndexDomainWrapper) = pyconvert(Int, parent(w).rank) function Base.axes(w::IndexDomainWrapper)