comparison src/LazyTensors/lazy_tensor_operations.jl @ 1001:a3df203861d3 refactor/lazy_tensors

Simplify some type parameter usage
author Jonatan Werpers <jonatan@werpers.com>
date Fri, 18 Mar 2022 22:01:25 +0100
parents 1091ac8c69ad
children 271aa6ae1055
comparison
equal deleted inserted replaced
1000:1091ac8c69ad 1001:a3df203861d3
1 # TBD: Is there a good way to split this file? 1 # TBD: Is there a good way to split this file?
2 # TODO: Split out functions for composition 2 # TODO: Split out functions for composition
3 # TODO: We need to be really careful about good error messages. 3 # TODO: We need to be really careful about good error messages.
4 # TODO: Go over type parameters
4 5
5 """ 6 """
6 LazyTensorApplication{T,R,D} <: LazyArray{T,R} 7 LazyTensorApplication{T,R,D} <: LazyArray{T,R}
7 8
8 Struct for lazy application of a LazyTensor. Created using `*`. 9 Struct for lazy application of a LazyTensor. Created using `*`.
56 # TODO: Rename this 57 # TODO: Rename this
57 struct LazyTensorBinaryOperation{Op,T,R,D,T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}} <: LazyTensor{T,D,R} 58 struct LazyTensorBinaryOperation{Op,T,R,D,T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}} <: LazyTensor{T,D,R}
58 tm1::T1 59 tm1::T1
59 tm2::T2 60 tm2::T2
60 61
61 @inline function LazyTensorBinaryOperation{Op,T,R,D}(tm1::T1,tm2::T2) where {Op,T,R,D, T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}} 62 function LazyTensorBinaryOperation{Op,T,R,D}(tm1::T1,tm2::T2) where {Op,T,R,D, T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}}
62 return new{Op,T,R,D,T1,T2}(tm1,tm2) 63 return new{Op,T,R,D,T1,T2}(tm1,tm2)
63 end 64 end
64 end 65 end
65 # TODO: Boundschecking in constructor. 66 # TODO: Boundschecking in constructor.
67
68 LazyTensorBinaryOperation{Op}(s,t) where Op = LazyTensorBinaryOperation{Op,eltype(s), range_dim(s), domain_dim(s)}(s,t)
66 69
67 apply(tmBinOp::LazyTensorBinaryOperation{:+,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) + apply(tmBinOp.tm2, v, I...) 70 apply(tmBinOp::LazyTensorBinaryOperation{:+,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) + apply(tmBinOp.tm2, v, I...)
68 apply(tmBinOp::LazyTensorBinaryOperation{:-,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) - apply(tmBinOp.tm2, v, I...) 71 apply(tmBinOp::LazyTensorBinaryOperation{:-,T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D} = apply(tmBinOp.tm1, v, I...) - apply(tmBinOp.tm2, v, I...)
69 72
70 range_size(tmBinOp::LazyTensorBinaryOperation) = range_size(tmBinOp.tm1) 73 range_size(tmBinOp::LazyTensorBinaryOperation) = range_size(tmBinOp.tm1)
78 """ 81 """
79 struct LazyTensorComposition{T,R,K,D, TM1<:LazyTensor{T,R,K}, TM2<:LazyTensor{T,K,D}} <: LazyTensor{T,R,D} 82 struct LazyTensorComposition{T,R,K,D, TM1<:LazyTensor{T,R,K}, TM2<:LazyTensor{T,K,D}} <: LazyTensor{T,R,D}
80 t1::TM1 83 t1::TM1
81 t2::TM2 84 t2::TM2
82 85
83 @inline function LazyTensorComposition(t1::LazyTensor{T,R,K}, t2::LazyTensor{T,K,D}) where {T,R,K,D} 86 function LazyTensorComposition(t1::LazyTensor{T,R,K}, t2::LazyTensor{T,K,D}) where {T,R,K,D}
84 @boundscheck check_domain_size(t1, range_size(t2)) 87 @boundscheck check_domain_size(t1, range_size(t2))
85 return new{T,R,K,D, typeof(t1), typeof(t2)}(t1,t2) 88 return new{T,R,K,D, typeof(t1), typeof(t2)}(t1,t2)
86 end 89 end
87 end 90 end
88 91