--- a/src/LazyTensors/LazyTensors.jl	Fri Mar 18 22:01:25 2022 +0100
+++ b/src/LazyTensors/LazyTensors.jl	Fri Mar 18 22:18:04 2022 +0100
@@ -12,6 +12,7 @@
 export SizeMismatch
 
 include("lazy_tensor.jl")
+include("tensor_types.jl")
 include("lazy_array.jl")
 include("lazy_tensor_operations.jl")
 include("tuple_manipulation.jl")

--- a/src/LazyTensors/lazy_tensor_operations.jl	Fri Mar 18 22:01:25 2022 +0100
+++ b/src/LazyTensors/lazy_tensor_operations.jl	Fri Mar 18 22:18:04 2022 +0100
@@ -1,8 +1,7 @@
-# TBD: Is there a good way to split this file?
-# TODO: Split out functions for composition
 # TODO: We need to be really careful about good error messages.
 # TODO: Go over type parameters
 
+
 """
     LazyTensorApplication{T,R,D} <: LazyArray{T,R}
 
@@ -54,7 +53,7 @@
 range_size(tmt::LazyTensorTranspose) = domain_size(tmt.tm)
 domain_size(tmt::LazyTensorTranspose) = range_size(tmt.tm)
 
-# TODO: Rename this
+
 struct LazyTensorBinaryOperation{Op,T,R,D,T1<:LazyTensor{T,R,D},T2<:LazyTensor{T,R,D}} <: LazyTensor{T,D,R}
     tm1::T1
     tm2::T2
@@ -102,66 +101,6 @@
 
 
 """
-    LazyLinearMap{T,R,D,...}(A, range_indicies, domain_indicies)
-
-LazyTensor defined by the AbstractArray A. `range_indicies` and `domain_indicies` define which indicies of A should
-be considerd the range and domain of the LazyTensor. Each set of indices must be ordered in ascending order.
-
-For instance, if A is a m x n matrix, and range_size = (1,), domain_size = (2,), then the LazyLinearMap performs the
-standard matrix-vector product on vectors of size n.
-"""
-struct LazyLinearMap{T,R,D, RD, AA<:AbstractArray{T,RD}} <: LazyTensor{T,R,D}
-    A::AA
-    range_indicies::NTuple{R,Int}
-    domain_indicies::NTuple{D,Int}
-
-    function LazyLinearMap(A::AA, range_indicies::NTuple{R,Int}, domain_indicies::NTuple{D,Int}) where {T,R,D, RD, AA<:AbstractArray{T,RD}}
-        if !issorted(range_indicies) || !issorted(domain_indicies)
-            throw(DomainError("range_indicies and domain_indicies must be sorted in ascending order"))
-        end
-
-        return new{T,R,D,RD,AA}(A,range_indicies,domain_indicies)
-    end
-end
-
-range_size(llm::LazyLinearMap) = size(llm.A)[[llm.range_indicies...]]
-domain_size(llm::LazyLinearMap) = size(llm.A)[[llm.domain_indicies...]]
-
-function apply(llm::LazyLinearMap{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D}
-    view_index = ntuple(i->:,ndims(llm.A))
-    for i ∈ 1:R
-        view_index = Base.setindex(view_index, Int(I[i]), llm.range_indicies[i])
-    end
-    A_view = @view llm.A[view_index...]
-    return sum(A_view.*v)
-end
-
-function apply_transpose(llm::LazyLinearMap{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D}
-    apply(LazyLinearMap(llm.A, llm.domain_indicies, llm.range_indicies), v, I...)
-end
-
-
-"""
-    IdentityTensor{T,D} <: LazyTensor{T,D,D}
-
-The lazy identity LazyTensor for a given size. Usefull for building up higher dimensional tensor mappings from lower
-dimensional ones through outer products. Also used in the Implementation for InflatedLazyTensor.
-"""
-struct IdentityTensor{T,D} <: LazyTensor{T,D,D}
-    size::NTuple{D,Int}
-end
-
-IdentityTensor{T}(size::NTuple{D,Int}) where {T,D} = IdentityTensor{T,D}(size)
-IdentityTensor{T}(size::Vararg{Int,D}) where {T,D} = IdentityTensor{T,D}(size)
-IdentityTensor(size::Vararg{Int,D}) where D = IdentityTensor{Float64,D}(size)
-
-range_size(tmi::IdentityTensor) = tmi.size
-domain_size(tmi::IdentityTensor) = tmi.size
-
-apply(tmi::IdentityTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = v[I...]
-apply_transpose(tmi::IdentityTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = v[I...]
-
-"""
     LazyTensorComposition(tm, tmi::IdentityTensor)
     LazyTensorComposition(tmi::IdentityTensor, tm)
 
@@ -182,21 +121,6 @@
     return tmi
 end
 
-"""
-    ScalingTensor{T,D} <: LazyTensor{T,D,D}
-
-A lazy tensor that scales its input with `λ`.
-"""
-struct ScalingTensor{T,D} <: LazyTensor{T,D,D}
-    λ::T
-    size::NTuple{D,Int}
-end
-
-LazyTensors.apply(tm::ScalingTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = tm.λ*v[I...]
-LazyTensors.apply_transpose(tm::ScalingTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = tm.λ*v[I...]
-
-LazyTensors.range_size(m::ScalingTensor) = m.size
-LazyTensors.domain_size(m::ScalingTensor) = m.size
 
 """
     InflatedLazyTensor{T,R,D} <: LazyTensor{T,R,D}
@@ -221,6 +145,7 @@
         return new{T,R,D,D_before,R_middle,D_middle,D_after, typeof(tm)}(before, tm, after)
     end
 end
+
 """
     InflatedLazyTensor(before, tm, after)
     InflatedLazyTensor(before,tm)

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/src/LazyTensors/tensor_types.jl	Fri Mar 18 22:18:04 2022 +0100
@@ -0,0 +1,76 @@
+"""
+    IdentityTensor{T,D} <: LazyTensor{T,D,D}
+
+The lazy identity LazyTensor for a given size. Usefull for building up higher dimensional tensor mappings from lower
+dimensional ones through outer products. Also used in the Implementation for InflatedLazyTensor.
+"""
+struct IdentityTensor{T,D} <: LazyTensor{T,D,D}
+    size::NTuple{D,Int}
+end
+
+IdentityTensor{T}(size::NTuple{D,Int}) where {T,D} = IdentityTensor{T,D}(size)
+IdentityTensor{T}(size::Vararg{Int,D}) where {T,D} = IdentityTensor{T,D}(size)
+IdentityTensor(size::Vararg{Int,D}) where D = IdentityTensor{Float64,D}(size)
+
+range_size(tmi::IdentityTensor) = tmi.size
+domain_size(tmi::IdentityTensor) = tmi.size
+
+apply(tmi::IdentityTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = v[I...]
+apply_transpose(tmi::IdentityTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = v[I...]
+
+
+"""
+    ScalingTensor{T,D} <: LazyTensor{T,D,D}
+
+A lazy tensor that scales its input with `λ`.
+"""
+struct ScalingTensor{T,D} <: LazyTensor{T,D,D}
+    λ::T
+    size::NTuple{D,Int}
+end
+
+LazyTensors.apply(tm::ScalingTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = tm.λ*v[I...]
+LazyTensors.apply_transpose(tm::ScalingTensor{T,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,D}) where {T,D} = tm.λ*v[I...]
+
+LazyTensors.range_size(m::ScalingTensor) = m.size
+LazyTensors.domain_size(m::ScalingTensor) = m.size
+
+
+"""
+    LazyLinearMap{T,R,D,...}(A, range_indicies, domain_indicies)
+
+LazyTensor defined by the AbstractArray A. `range_indicies` and `domain_indicies` define which indicies of A should
+be considerd the range and domain of the LazyTensor. Each set of indices must be ordered in ascending order.
+
+For instance, if A is a m x n matrix, and range_size = (1,), domain_size = (2,), then the LazyLinearMap performs the
+standard matrix-vector product on vectors of size n.
+"""
+struct LazyLinearMap{T,R,D, RD, AA<:AbstractArray{T,RD}} <: LazyTensor{T,R,D}
+    A::AA
+    range_indicies::NTuple{R,Int}
+    domain_indicies::NTuple{D,Int}
+
+    function LazyLinearMap(A::AA, range_indicies::NTuple{R,Int}, domain_indicies::NTuple{D,Int}) where {T,R,D, RD, AA<:AbstractArray{T,RD}}
+        if !issorted(range_indicies) || !issorted(domain_indicies)
+            throw(DomainError("range_indicies and domain_indicies must be sorted in ascending order"))
+        end
+
+        return new{T,R,D,RD,AA}(A,range_indicies,domain_indicies)
+    end
+end
+
+range_size(llm::LazyLinearMap) = size(llm.A)[[llm.range_indicies...]]
+domain_size(llm::LazyLinearMap) = size(llm.A)[[llm.domain_indicies...]]
+
+function apply(llm::LazyLinearMap{T,R,D}, v::AbstractArray{<:Any,D}, I::Vararg{Any,R}) where {T,R,D}
+    view_index = ntuple(i->:,ndims(llm.A))
+    for i ∈ 1:R
+        view_index = Base.setindex(view_index, Int(I[i]), llm.range_indicies[i])
+    end
+    A_view = @view llm.A[view_index...]
+    return sum(A_view.*v)
+end
+
+function apply_transpose(llm::LazyLinearMap{T,R,D}, v::AbstractArray{<:Any,R}, I::Vararg{Any,D}) where {T,R,D}
+    apply(LazyLinearMap(llm.A, llm.domain_indicies, llm.range_indicies), v, I...)
+end

--- a/test/LazyTensors/lazy_tensor_operations_test.jl	Fri Mar 18 22:01:25 2022 +0100
+++ b/test/LazyTensors/lazy_tensor_operations_test.jl	Fri Mar 18 22:18:04 2022 +0100
@@ -24,6 +24,7 @@
     @test domain_size(m') == :range_size
 end
 
+
 @testset "TensorApplication" begin
     struct SizeDoublingMapping{T,R,D} <: LazyTensor{T,R,D}
         domain_size::NTuple{D,Int}
@@ -108,6 +109,7 @@
     end
 end
 
+
 @testset "LazyTensor binary operations" begin
     A = ScalingTensor(2.0, (3,))
     B = ScalingTensor(3.0, (3,))
@@ -155,111 +157,6 @@
     @test ((Ã∘B̃)'*ComplexF64[1.,2.])[1] isa ComplexF64
 end
 
-@testset "LazyLinearMap" begin
-    # Test a standard matrix-vector product
-    # mapping vectors of size 4 to vectors of size 3.
-    A = rand(3,4)
-    Ã = LazyLinearMap(A, (1,), (2,))
-    v = rand(4)
-    w = rand(3)
-
-    @test à isa LazyLinearMap{T,1,1} where T
-    @test à isa LazyTensor{T,1,1} where T
-    @test range_size(Ã) == (3,)
-    @test domain_size(Ã) == (4,)
-
-    @test Ã*ones(4) ≈ A*ones(4) atol=5e-13
-    @test Ã*v ≈ A*v atol=5e-13
-    @test Ã'*w ≈ A'*w
-
-    A = rand(2,3,4)
-    @test_throws DomainError LazyLinearMap(A, (3,1), (2,))
-
-    # Test more exotic mappings
-    B = rand(3,4,2)
-    # Map vectors of size 2 to matrices of size (3,4)
-    B̃ = LazyLinearMap(B, (1,2), (3,))
-    v = rand(2)
-
-    @test range_size(B̃) == (3,4)
-    @test domain_size(B̃) == (2,)
-    @test B̃ isa LazyTensor{T,2,1} where T
-    @test B̃*ones(2) ≈ B[:,:,1] + B[:,:,2] atol=5e-13
-    @test B̃*v ≈ B[:,:,1]*v[1] + B[:,:,2]*v[2] atol=5e-13
-
-    # Map matrices of size (3,2) to vectors of size 4
-    B̃ = LazyLinearMap(B, (2,), (1,3))
-    v = rand(3,2)
-
-    @test range_size(B̃) == (4,)
-    @test domain_size(B̃) == (3,2)
-    @test B̃ isa LazyTensor{T,1,2} where T
-    @test B̃*ones(3,2) ≈ B[1,:,1] + B[2,:,1] + B[3,:,1] +
-                        B[1,:,2] + B[2,:,2] + B[3,:,2] atol=5e-13
-    @test B̃*v ≈ B[1,:,1]*v[1,1] + B[2,:,1]*v[2,1] + B[3,:,1]*v[3,1] +
-                B[1,:,2]v[1,2] + B[2,:,2]*v[2,2] + B[3,:,2]*v[3,2] atol=5e-13
-
-
-    # TODO:
-    # @inferred (B̃*v)[2]
-end
-
-
-@testset "IdentityTensor" begin
-    @test IdentityTensor{Float64}((4,5)) isa IdentityTensor{T,2} where T
-    @test IdentityTensor{Float64}((4,5)) isa LazyTensor{T,2,2} where T
-    @test IdentityTensor{Float64}((4,5)) == IdentityTensor{Float64}(4,5)
-
-    @test IdentityTensor(3,2) isa IdentityTensor{Float64,2}
-
-    for sz ∈ [(4,5),(3,),(5,6,4)]
-        I = IdentityTensor{Float64}(sz)
-        v = rand(sz...)
-        @test I*v == v
-        @test I'*v == v
-
-        v = rand(ComplexF64,sz...)
-        @test I*v == v
-        @test I'*v == v
-
-        @test range_size(I) == sz
-        @test domain_size(I) == sz
-    end
-
-    I = IdentityTensor{Float64}((4,5))
-    v = rand(4,5)
-    @inferred (I*v)[3,2]
-    @inferred (I'*v)[3,2]
-    @inferred range_size(I)
-
-    @inferred range_dim(I)
-    @inferred domain_dim(I)
-
-    Ã = rand(4,2)
-    A = LazyLinearMap(Ã,(1,),(2,))
-    I1 = IdentityTensor{Float64}(2)
-    I2 = IdentityTensor{Float64}(4)
-    @test A∘I1 == A
-    @test I2∘A == A
-    @test I1∘I1 == I1
-    @test_throws SizeMismatch I1∘A
-    @test_throws SizeMismatch A∘I2
-    @test_throws SizeMismatch I1∘I2
-end
-
-@testset "ScalingTensor" begin
-    st = ScalingTensor(2.,(3,4))
-    @test st isa LazyTensor{Float64, 2, 2}
-    @test range_size(st) == (3,4)
-    @test domain_size(st) == (3,4)
-
-    v = rand(3,4)
-    @test st*v == 2.0 .* v
-    @test st'*v == 2.0 .* v
-
-    @inferred (st*v)[2,2]
-    @inferred (st'*v)[2,2]
-end
 
 @testset "InflatedLazyTensor" begin
     I(sz...) = IdentityTensor(sz...)
@@ -395,7 +292,6 @@
 end
 
 @testset "LazyOuterProduct" begin
-
     A = ScalingTensor(2.0, (5,))
     B = ScalingTensor(3.0, (3,))
     C = ScalingTensor(5.0, (3,2))
@@ -445,5 +341,4 @@
         I2 = IdentityTensor(4)
         @test I1⊗Ã⊗I2 == InflatedLazyTensor(I1, Ã, I2)
     end
-
 end

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/test/LazyTensors/tensor_types_test.jl	Fri Mar 18 22:18:04 2022 +0100
@@ -0,0 +1,109 @@
+using Test
+using Sbplib.LazyTensors
+
+@testset "IdentityTensor" begin
+    @test IdentityTensor{Float64}((4,5)) isa IdentityTensor{T,2} where T
+    @test IdentityTensor{Float64}((4,5)) isa LazyTensor{T,2,2} where T
+    @test IdentityTensor{Float64}((4,5)) == IdentityTensor{Float64}(4,5)
+
+    @test IdentityTensor(3,2) isa IdentityTensor{Float64,2}
+
+    for sz ∈ [(4,5),(3,),(5,6,4)]
+        I = IdentityTensor{Float64}(sz)
+        v = rand(sz...)
+        @test I*v == v
+        @test I'*v == v
+
+        v = rand(ComplexF64,sz...)
+        @test I*v == v
+        @test I'*v == v
+
+        @test range_size(I) == sz
+        @test domain_size(I) == sz
+    end
+
+    I = IdentityTensor{Float64}((4,5))
+    v = rand(4,5)
+    @inferred (I*v)[3,2]
+    @inferred (I'*v)[3,2]
+    @inferred range_size(I)
+
+    @inferred range_dim(I)
+    @inferred domain_dim(I)
+
+    Ã = rand(4,2)
+    A = LazyLinearMap(Ã,(1,),(2,))
+    I1 = IdentityTensor{Float64}(2)
+    I2 = IdentityTensor{Float64}(4)
+    @test A∘I1 == A
+    @test I2∘A == A
+    @test I1∘I1 == I1
+    @test_throws SizeMismatch I1∘A
+    @test_throws SizeMismatch A∘I2
+    @test_throws SizeMismatch I1∘I2
+end
+
+
+@testset "ScalingTensor" begin
+    st = ScalingTensor(2.,(3,4))
+    @test st isa LazyTensor{Float64, 2, 2}
+    @test range_size(st) == (3,4)
+    @test domain_size(st) == (3,4)
+
+    v = rand(3,4)
+    @test st*v == 2.0 .* v
+    @test st'*v == 2.0 .* v
+
+    @inferred (st*v)[2,2]
+    @inferred (st'*v)[2,2]
+end
+
+
+@testset "LazyLinearMap" begin
+    # Test a standard matrix-vector product
+    # mapping vectors of size 4 to vectors of size 3.
+    A = rand(3,4)
+    Ã = LazyLinearMap(A, (1,), (2,))
+    v = rand(4)
+    w = rand(3)
+
+    @test à isa LazyLinearMap{T,1,1} where T
+    @test à isa LazyTensor{T,1,1} where T
+    @test range_size(Ã) == (3,)
+    @test domain_size(Ã) == (4,)
+
+    @test Ã*ones(4) ≈ A*ones(4) atol=5e-13
+    @test Ã*v ≈ A*v atol=5e-13
+    @test Ã'*w ≈ A'*w
+
+    A = rand(2,3,4)
+    @test_throws DomainError LazyLinearMap(A, (3,1), (2,))
+
+    # Test more exotic mappings
+    B = rand(3,4,2)
+    # Map vectors of size 2 to matrices of size (3,4)
+    B̃ = LazyLinearMap(B, (1,2), (3,))
+    v = rand(2)
+
+    @test range_size(B̃) == (3,4)
+    @test domain_size(B̃) == (2,)
+    @test B̃ isa LazyTensor{T,2,1} where T
+    @test B̃*ones(2) ≈ B[:,:,1] + B[:,:,2] atol=5e-13
+    @test B̃*v ≈ B[:,:,1]*v[1] + B[:,:,2]*v[2] atol=5e-13
+
+    # Map matrices of size (3,2) to vectors of size 4
+    B̃ = LazyLinearMap(B, (2,), (1,3))
+    v = rand(3,2)
+
+    @test range_size(B̃) == (4,)
+    @test domain_size(B̃) == (3,2)
+    @test B̃ isa LazyTensor{T,1,2} where T
+    @test B̃*ones(3,2) ≈ B[1,:,1] + B[2,:,1] + B[3,:,1] +
+                        B[1,:,2] + B[2,:,2] + B[3,:,2] atol=5e-13
+    @test B̃*v ≈ B[1,:,1]*v[1,1] + B[2,:,1]*v[2,1] + B[3,:,1]*v[3,1] +
+                B[1,:,2]v[1,2] + B[2,:,2]*v[2,2] + B[3,:,2]*v[3,2] atol=5e-13
+
+
+    # TODO:
+    # @inferred (B̃*v)[2]
+end

--- a/test/LazyTensors/tuple_manipulation_test.jl	Fri Mar 18 22:01:25 2022 +0100
+++ b/test/LazyTensors/tuple_manipulation_test.jl	Fri Mar 18 22:18:04 2022 +0100
@@ -1,3 +1,6 @@
+using Test
+using Sbplib.LazyTensors
+
 @testset "split_index" begin
     @test LazyTensors.split_index(Val(2),Val(1),Val(2),Val(2),1,2,3,4,5,6) == ((1,2,:,5,6),(3,4))
     @test LazyTensors.split_index(Val(2),Val(3),Val(2),Val(2),1,2,3,4,5,6) == ((1,2,:,:,:,5,6),(3,4))