changeset 1012:2c1a0722ddb9 feature/lazy_tensors/pretty_printing

Add pretty printing for inflated tensor
author Jonatan Werpers <jonatan@werpers.com>
date Mon, 21 Mar 2022 09:07:28 +0100
parents f2c5d44a8db0
children 67969bd7e642
files src/LazyTensors/lazy_tensor_operations.jl test/LazyTensors/lazy_tensor_operations_test.jl
diffstat 2 files changed, 23 insertions(+), 0 deletions(-) [+]
line wrap: on
line diff
--- a/src/LazyTensors/lazy_tensor_operations.jl	Mon Mar 21 08:55:54 2022 +0100
+++ b/src/LazyTensors/lazy_tensor_operations.jl	Mon Mar 21 09:07:28 2022 +0100
@@ -214,6 +214,15 @@
     return apply_transpose(itm.tm, v_inner, inner_index...)
 end
 
+function Base.show(io::IO, ::MIME"text/plain", tm::InflatedLazyTensor{T}) where T
+    show(IOContext(io, :compact=>true), MIME("text/plain"), tm.before)
+    print(io, "⊗")
+    # if get(io, :compact, false)
+    show(io, MIME("text/plain"), tm.tm)
+    print(io, "⊗")
+    show(IOContext(io, :compact=>true), MIME("text/plain"), tm.after)
+end
+
 
 @doc raw"""
     LazyOuterProduct(tms...)
--- a/test/LazyTensors/lazy_tensor_operations_test.jl	Mon Mar 21 08:55:54 2022 +0100
+++ b/test/LazyTensors/lazy_tensor_operations_test.jl	Mon Mar 21 09:07:28 2022 +0100
@@ -315,6 +315,20 @@
 
         @test InflatedLazyTensor(I(2), I(2), I(2)) isa InflatedLazyTensor # The constructor should always return its type.
     end
+
+    @testset "Pretty printing" begin
+        cases = [
+            InflatedLazyTensor(I(4), ScalingTensor(2., (3,2)), I(2)) => (
+                regular="I(4)⊗ScalingTensor{Float64}(2.0, (3, 2))⊗I(2)",
+                compact="I(4)⊗2.0*I(3,2)⊗I(2)"
+            )
+        ]
+
+        @testset "$tm" for (tm, r) ∈ cases
+            @test repr(MIME("text/plain"), tm) == r.regular
+            @test repr(MIME("text/plain"), tm, context=:compact=>true) == r.compact
+        end
+    end
 end
 
 @testset "LazyOuterProduct" begin