From a5cab3ea11edccb1ee0976356e7fa80e74db3941 Mon Sep 17 00:00:00 2001
From: rafaqz <Rafael Schouten>
Date: Sun, 29 Aug 2021 00:34:43 +0200
Subject: [PATCH] throw DimensionMismatch for broadcast on mixed dimensions

---
 src/ops.jl       | 11 ++++++++---
 test/runtests.jl |  6 ++++++
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/ops.jl b/src/ops.jl
index c95fc05..86dcd18 100644
--- a/src/ops.jl
+++ b/src/ops.jl
@@ -62,9 +62,14 @@ function common_chunks(s,args...)
     return map(i->i[1],tt),map(i->i[2],tt)
   end
 end
-subsetarg(x, a) = x
-function subsetarg(x::AbstractArray,a)
-  ashort = maybeonerange(size(x),a)
+subsetarg(x, ranges) = x
+function subsetarg(x::Tuple, ranges)
+  length(x) !== 1 && length(ranges) !== 1 && throw(DimensionMismatch("DiskArrays can only broadcast a Tuple of length > 1 with a 1 dimensional Array"))
+  x
+end
+function subsetarg(x::AbstractArray,ranges)
+  (length(ranges) === ndims(x) || length(x) === 1) || throw(DimensionMismatch("Can only broadcast with arrays of matching size"))
+  ashort = maybeonerange(size(x),ranges)
   view(x,ashort...) #Maybe making a copy here would be faster, need to check...
 end
 repsingle(s,r) = s==1 ? (1:1) : r
diff --git a/test/runtests.jl b/test/runtests.jl
index fb316fd..e2c3e64 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -207,6 +207,12 @@ end
 @testset "Broadcast" begin
   a_disk1 = _DiskArray(rand(10,9,2), chunksize=(5,3,2))
   test_broadcast(a_disk1)
+  a = ones(10,9,2);
+  a_disk = _DiskArray(a, chunksize=(5,3,2));
+  tup = ntuple(x -> x, 10)
+  mat = reshape(1:90, 10, 9)
+  @test_throws DimensionMismatch a_disk .* tup |> collect
+  @test_throws DimensionMismatch a_disk .* mat |> collect
 end
 
 @testset "Broadcast with length 1 final dim" begin