diff src/LazyTensors/lazy_tensor_operations.jl @ 417:4c6604b7d990 feature/tensor_composition

Add dimension checking in the constructor
author Jonatan Werpers <jonatan@werpers.com>
date Fri, 16 Oct 2020 20:32:09 +0200
parents 814865d40f48
children 264af2bb646f
line wrap: on
line diff
--- a/src/LazyTensors/lazy_tensor_operations.jl	Fri Oct 16 09:37:35 2020 +0200
+++ b/src/LazyTensors/lazy_tensor_operations.jl	Fri Oct 16 20:32:09 2020 +0200
@@ -83,6 +83,12 @@
     t1::TM1
     t2::TM2
 
+    @inline function TensorMappingComposition(t1::TensorMapping{T,R,K}, t2::TensorMapping{T,K,D}) where {T,R,K,D}
+        @boundscheck if domain_size(t1) != range_size(t2)
+            throw(DimensionMismatch("The first argument has domain size $(domain_size(t1)) while the second has range size $(range_size(t2)) "))
+        end
+        return new{T,R,K,D, typeof(t1), typeof(t2)}(t1,t2)
+    end
     # Add check for matching sizes as a boundscheck
 end
 export TensorMappingComposition
@@ -98,7 +104,7 @@
     apply_transpose(c.t2, LazyTensorMappingApplication(c.t1',v), I...)
 end
 
-Base.:∘(s::TensorMapping, t::TensorMapping) = TensorMappingComposition(s,t)
+Base.@propagate_inbounds Base.:∘(s::TensorMapping, t::TensorMapping) = TensorMappingComposition(s,t)
 
 """
     LazyLinearMap{T,R,D,...}(A, range_indicies, domain_indicies)