comparison src/LazyTensors/lazy_tensor_operations.jl @ 348:7fe43d902a27 refactor/remove_dynamic_size_tensormapping

Start trying to change LazyTensors
author Jonatan Werpers <jonatan@werpers.com>
date Sat, 26 Sep 2020 15:20:18 +0200
parents 01b851161018
children ba46a952a450
comparison
equal deleted inserted replaced
344:f781d6da7d3d 348:7fe43d902a27
9 """ 9 """
10 struct LazyTensorMappingApplication{T,R,D} <: LazyArray{T,R} 10 struct LazyTensorMappingApplication{T,R,D} <: LazyArray{T,R}
11 t::TensorMapping{T,R,D} 11 t::TensorMapping{T,R,D}
12 o::AbstractArray{T,D} 12 o::AbstractArray{T,D}
13 end 13 end
14 # TODO: Do boundschecking on creation!
14 export LazyTensorMappingApplication 15 export LazyTensorMappingApplication
15 16
16 Base.:*(tm::TensorMapping{T,R,D}, o::AbstractArray{T,D}) where {T,R,D} = LazyTensorMappingApplication(tm,o) 17 Base.:*(tm::TensorMapping{T,R,D}, o::AbstractArray{T,D}) where {T,R,D} = LazyTensorMappingApplication(tm,o)
17 Base.getindex(ta::LazyTensorMappingApplication{T,R,D}, I::Vararg{Index,R}) where {T,R,D} = apply(ta.t, ta.o, I...) 18 Base.getindex(ta::LazyTensorMappingApplication{T,R,D}, I::Vararg{Index,R}) where {T,R,D} = apply(ta.t, ta.o, I...)
18 Base.getindex(ta::LazyTensorMappingApplication{T,R,D}, I::Vararg{Int,R}) where {T,R,D} = apply(ta.t, ta.o, Index{Unknown}.(I)...) 19 Base.getindex(ta::LazyTensorMappingApplication{T,R,D}, I::Vararg{Int,R}) where {T,R,D} = apply(ta.t, ta.o, Index{Unknown}.(I)...)
19 Base.size(ta::LazyTensorMappingApplication{T,R,D}) where {T,R,D} = range_size(ta.t,size(ta.o)) 20 Base.size(ta::LazyTensorMappingApplication{T,R,D}) where {T,R,D} = range_size(ta.t)
20 # TODO: What else is needed to implement the AbstractArray interface? 21 # TODO: What else is needed to implement the AbstractArray interface?
21 22
22 # # We need the associativity to be a→b→c = a→(b→c), which is the case for '→' 23 # # We need the associativity to be a→b→c = a→(b→c), which is the case for '→'
23 Base.:*(a::TensorMapping{T,R,D}, b::TensorMapping{T,D,K}, args::Union{TensorMapping{T}, AbstractArray{T}}...) where {T,R,D,K} = foldr(*,(a,b,args...)) 24 Base.:*(a::TensorMapping{T,R,D}, b::TensorMapping{T,D,K}, args::Union{TensorMapping{T}, AbstractArray{T}}...) where {T,R,D,K} = foldr(*,(a,b,args...))
24 # # Should we overload some other infix binary opesrator? 25 # # Should we overload some other infix binary opesrator?
39 tm::TensorMapping{T,R,D} 40 tm::TensorMapping{T,R,D}
40 end 41 end
41 export LazyTensorMappingTranspose 42 export LazyTensorMappingTranspose
42 43
43 # # TBD: Should this be implemented on a type by type basis or through a trait to provide earlier errors? 44 # # TBD: Should this be implemented on a type by type basis or through a trait to provide earlier errors?
45 # Jonatan 2020-09-25: Is the problem that you can take the transpose of any TensorMapping even if it doesn't implement `apply_transpose`?
44 Base.adjoint(tm::TensorMapping) = LazyTensorMappingTranspose(tm) 46 Base.adjoint(tm::TensorMapping) = LazyTensorMappingTranspose(tm)
45 Base.adjoint(tmt::LazyTensorMappingTranspose) = tmt.tm 47 Base.adjoint(tmt::LazyTensorMappingTranspose) = tmt.tm
46 48
47 apply(tmt::LazyTensorMappingTranspose{T,R,D}, v::AbstractArray{T,R}, I::Vararg{Index,D}) where {T,R,D} = apply_transpose(tmt.tm, v, I...) 49 apply(tmt::LazyTensorMappingTranspose{T,R,D}, v::AbstractArray{T,R}, I::Vararg{Index,D}) where {T,R,D} = apply_transpose(tmt.tm, v, I...)
48 apply_transpose(tmt::LazyTensorMappingTranspose{T,R,D}, v::AbstractArray{T,D}, I::Vararg{Index,R}) where {T,R,D} = apply(tmt.tm, v, I...) 50 apply_transpose(tmt::LazyTensorMappingTranspose{T,R,D}, v::AbstractArray{T,D}, I::Vararg{Index,R}) where {T,R,D} = apply(tmt.tm, v, I...)
49 51
50 range_size(tmt::LazyTensorMappingTranspose{T,R,D}, d_size::NTuple{R,Integer}) where {T,R,D} = domain_size(tmt.tm, d_size) 52 range_size(tmt::LazyTensorMappingTranspose{T,R,D}) where {T,R,D} = domain_size(tmt.tm)
51 domain_size(tmt::LazyTensorMappingTranspose{T,R,D}, r_size::NTuple{D,Integer}) where {T,R,D} = range_size(tmt.tm, r_size) 53 domain_size(tmt::LazyTensorMappingTranspose{T,R,D}) where {T,R,D} = range_size(tmt.tm)
52
53
54 54
55 55
56 struct LazyTensorMappingBinaryOperation{Op,T,R,D,T1<:TensorMapping{T,R,D},T2<:TensorMapping{T,R,D}} <: TensorMapping{T,D,R} 56 struct LazyTensorMappingBinaryOperation{Op,T,R,D,T1<:TensorMapping{T,R,D},T2<:TensorMapping{T,R,D}} <: TensorMapping{T,D,R}
57 tm1::T1 57 tm1::T1
58 tm2::T2 58 tm2::T2
59 59
60 @inline function LazyTensorMappingBinaryOperation{Op,T,R,D}(tm1::T1,tm2::T2) where {Op,T,R,D, T1<:TensorMapping{T,R,D},T2<:TensorMapping{T,R,D}} 60 @inline function LazyTensorMappingBinaryOperation{Op,T,R,D}(tm1::T1,tm2::T2) where {Op,T,R,D, T1<:TensorMapping{T,R,D},T2<:TensorMapping{T,R,D}}
61 return new{Op,T,R,D,T1,T2}(tm1,tm2) 61 return new{Op,T,R,D,T1,T2}(tm1,tm2)
62 end 62 end
63 end 63 end
64 # TODO: Boundschecking in constructor.
64 65
65 apply(tmBinOp::LazyTensorMappingBinaryOperation{:+,T,R,D}, v::AbstractArray{T,D}, I::Vararg{Index,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) + apply(tmBinOp.tm2, v, I...) 66 apply(tmBinOp::LazyTensorMappingBinaryOperation{:+,T,R,D}, v::AbstractArray{T,D}, I::Vararg{Index,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) + apply(tmBinOp.tm2, v, I...)
66 apply(tmBinOp::LazyTensorMappingBinaryOperation{:-,T,R,D}, v::AbstractArray{T,D}, I::Vararg{Index,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) - apply(tmBinOp.tm2, v, I...) 67 apply(tmBinOp::LazyTensorMappingBinaryOperation{:-,T,R,D}, v::AbstractArray{T,D}, I::Vararg{Index,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) - apply(tmBinOp.tm2, v, I...)
67 68
68 range_size(tmBinOp::LazyTensorMappingBinaryOperation{Op,T,R,D}, domain_size::NTuple{D,Integer}) where {Op,T,R,D} = range_size(tmBinOp.tm1, domain_size) 69 range_size(tmBinOp::LazyTensorMappingBinaryOperation{Op,T,R,D}) where {Op,T,R,D} = range_size(tmBinOp.tm1)
69 domain_size(tmBinOp::LazyTensorMappingBinaryOperation{Op,T,R,D}, range_size::NTuple{R,Integer}) where {Op,T,R,D} = domain_size(tmBinOp.tm2, range_size) 70 domain_size(tmBinOp::LazyTensorMappingBinaryOperation{Op,T,R,D}) where {Op,T,R,D} = domain_size(tmBinOp.tm1)
70 71
71 Base.:+(tm1::TensorMapping{T,R,D}, tm2::TensorMapping{T,R,D}) where {T,R,D} = LazyTensorMappingBinaryOperation{:+,T,R,D}(tm1,tm2) 72 Base.:+(tm1::TensorMapping{T,R,D}, tm2::TensorMapping{T,R,D}) where {T,R,D} = LazyTensorMappingBinaryOperation{:+,T,R,D}(tm1,tm2)
72 Base.:-(tm1::TensorMapping{T,R,D}, tm2::TensorMapping{T,R,D}) where {T,R,D} = LazyTensorMappingBinaryOperation{:-,T,R,D}(tm1,tm2) 73 Base.:-(tm1::TensorMapping{T,R,D}, tm2::TensorMapping{T,R,D}) where {T,R,D} = LazyTensorMappingBinaryOperation{:-,T,R,D}(tm1,tm2)
73 74
74 75