put type text below matrices.

Author: parrt

testing/test2.py

@@ -7,8 +7,18 @@
 x = torch.tensor([4, 5], dtype=torch.int32).reshape(2, 1)
 h = torch.tensor([1,2])
 
-with tsensor.explain(legend=True, savefig="/Users/parrt/Desktop/t2.pdf"):
-    torch.rand(size=(2,20,2000,10))
+a = torch.rand(size=(2, 20), dtype=torch.float64)
+b = torch.rand(size=(2, 20), dtype=torch.float32)
+c = torch.rand(size=(2,20,200), dtype=torch.complex64)
+d = torch.rand(size=(2,20,200,5), dtype=torch.float16)
+with tsensor.explain(legend=False, savefig="/Users/parrt/Desktop/t2.pdf"):
+    a + b + x + c[:,:,0] + d[:,:,0,0]
+
+with tsensor.explain(legend=False, savefig="/Users/parrt/Desktop/t3.pdf"):
+    c
+
+with tsensor.explain(legend=False, savefig="/Users/parrt/Desktop/t4.pdf"):
+    d
 
 # with tsensor.explain(legend=True, savefig="/Users/parrt/Desktop/t.pdf") as e:
 #     W @ torch.dot(b, b) + torch.eye(2, 2) @ x

tsensor/viz.py

@@ -316,12 +316,20 @@ def draw_matrix(self,ax,sub, sh, ty):
                     fontname=self.dimfontname, fontsize=self.dimfontsize,
                     rotation=45)
 
+        bottom_text_line = self.box_topy - h - self.dim_ypadding
         if len(sh) > 3:
             # Text below
             remaining = r"$\cdots\mathsf{x}$"+r"$\mathsf{x}$".join([self.nabbrev(sh[i]) for i in range(3,len(sh))])
-            ax.text(mid, self.box_topy - h - self.dim_ypadding, remaining,
+            bottom_text_line = self.box_topy - h - self.dim_ypadding
+            ax.text(mid, bottom_text_line, remaining,
                     verticalalignment='top', horizontalalignment='center',
                     fontname=self.dimfontname, fontsize=self.dimfontsize)
+            bottom_text_line -= self.hchar
+
+        # Type info at the bottom of everything
+        ax.text(mid, bottom_text_line, '<${\mathit{'+ty+'}}$>',
+                verticalalignment='top', horizontalalignment='center',
+                fontname=self.dimfontname, fontsize=self.dimfontsize-2)
 
     @staticmethod
     def nabbrev(n: int) -> str: