Merge branch 'sbrugman-patch-2'

Author: parrt

testing/examples.ipynb

@@ -0,0 +1,53 @@
+import numpy as np
+
+import matplotlib.pyplot as plt
+import matplotlib.colors as mc
+
+bits = [4,8,16,32,64,128]
+bits = [8,16,32,64,128]
+nhues = len(bits)
+
+blueish = '#3B75AF'
+greenish = '#519E3E'
+
+orangeish = '#FDDB7D'
+limeish = '#C1E1C5'
+limeish = '#A8E1B0'
+yellowish = '#FFFFAD'
+
+print(mc.hex2color(limeish))
+
+type_colors = {'float':limeish, 'int':blueish, 'complex':orangeish}
+
+# Derived from https://stackoverflow.com/questions/47222585/matplotlib-generic-colormap-from-tab10
+
+def categorical_cmap(color, nsc):
+    # ccolors = plt.get_cmap(cmap)(np.arange(nc, dtype=int))
+    # print(ccolors[0:4])
+    cols = np.zeros((nsc, 3))
+    # chsv = mc.rgb_to_hsv(c[:3])
+    chsv = mc.rgb_to_hsv(mc.hex2color(color))
+    arhsv = np.tile(chsv,nsc).reshape(nsc,3)
+    arhsv[:,1] = np.linspace(chsv[1],0.25,nsc)
+    arhsv[:,2] = np.linspace(chsv[2],1,nsc)
+    rgb = mc.hsv_to_rgb(arhsv)
+    cols[0:nsc,:] = rgb
+    cmap = mc.ListedColormap(cols)
+    return cmap
+
+plt.figure(figsize=(3,3))
+c1 = categorical_cmap(blueish,nhues)
+plt.scatter(np.arange(nhues),[1]*nhues, c=np.arange(nhues), s=1080, cmap=c1, linewidths=.5, edgecolors='grey')
+c1 = categorical_cmap(limeish,nhues)
+plt.scatter(np.arange(nhues),[2]*nhues, c=np.arange(nhues), s=1080, cmap=c1, linewidths=.5, edgecolors='grey')
+c1 = categorical_cmap(yellowish,nhues)
+plt.scatter(np.arange(nhues),[3]*nhues, c=np.arange(nhues), s=1080, cmap=c1, linewidths=.5, edgecolors='grey')
+
+plt.margins(y=3)
+plt.xticks([])
+plt.yticks([0,1,2],["(5, 4)", "(2, 5)", "(4, 3)"])
+plt.ylim(0, 4)
+plt.axis('off')
+
+plt.savefig("/Users/parrt/Desktop/colors.pdf")
+plt.show()

testing/play_cmaps.py

@@ -1,16 +1,34 @@
-import torch
+import numpy as np
 import tsensor
-nhidden = 256
-n = 200         # number of instances
-d = 764         # number of instance features
-n_neurons = 100 # how many neurons in this layer?
-
-Whh_ = torch.eye(nhidden, nhidden)
-Uxh_ = torch.randn(nhidden, d)
-bh_  = torch.zeros(nhidden, 1)
-h = torch.randn(nhidden, 1)  # fake previous hidden state h
-r = torch.randn(nhidden, 1)  # fake this computation
-X = torch.rand(n,d)          # fake input
-
-g = tsensor.astviz("h_ = torch.tanh(Whh_ @ (r*h) + Uxh_ @ X.T + bh_)")
-g.savefig("/tmp/torch-gru-ast-shapes.svg")
+import torch
+
+W = np.array([[1, 2], [3, 4]])
+b = np.array([9, 10]).reshape(2, 1)
+x = np.array([4, 5]).reshape(2, 1)
+h = np.array([1, 2])
+with tsensor.explain(savefig="/Users/parrt/Desktop/foo.pdf"):
+    W @ np.dot(b,b) + np.eye(2,2)@x
+
+
+W = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
+b = torch.tensor([9, 10]).reshape(2, 1)
+x = torch.tensor([4, 5], dtype=torch.int32).reshape(2, 1)
+h = torch.tensor([1,2])
+
+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(savefig="/Users/parrt/Desktop/t2.pdf"):
+    a + b + x + c[:,:,0] + d[:,:,0,0]
+
+with tsensor.explain(savefig="/Users/parrt/Desktop/t3.pdf"):
+    c
+
+with tsensor.explain(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

testing/test2.py

@@ -1,19 +1,16 @@
-import tsensor
-import tensorflow as tf
 import matplotlib.pyplot as plt
+import numpy as np
+import tsensor
+import torch
 
-W = tf.constant([[1, 2], [3, 4]])
-b = tf.reshape(tf.constant([[9, 10]]), (2, 1))
-x = tf.reshape(tf.constant([[8, 5, 7]]), (3, 1))
-z = 0
-
-# tsensor.parse("z /= b + x * 3", hush_errors=False)
-
-# with tsensor.clarify(show='viz'):
-#     b + x * 3
+W = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
+b = torch.tensor([9, 10]).reshape(2, 1)
+x = torch.tensor([4, 5], dtype=torch.int32).reshape(2, 1)
+h = torch.tensor([1,2])
 
 fig, ax = plt.subplots(1,1)
-tsensor.pyviz("b + x", ax=ax)
+view = tsensor.pyviz("b + x", ax=ax, legend=True)
+# view.savefig("/Users/parrt/Desktop/foo.pdf")
 plt.show()
 # with tsensor.explain():
 #     b + x

testing/test3.py

@@ -0,0 +1,32 @@
+import numpy as np
+import jax.numpy as jnp
+import tensorflow as tf
+import torch
+import pytest
+
+import tsensor as ts
+
+
+@pytest.mark.parametrize(
+    "value,expected",
+    [
+        # Numpy
+        (np.random.randint(1, 10, size=(10, 2, 5)), "int64"),
+        (np.random.randint(1, 10, size=(10, 2, 5), dtype="int8"), "int8"),
+        (np.random.normal(size=(5, 1)).astype(np.float32), "float32"),
+        (np.random.normal(size=(5, 1)).astype(np.float32), "float32"),
+        (np.array([('Rex', 9, 81.0), ('Fido', 3, 27.0)], dtype=[('name', 'U10'), ('age', 'i4'), ('weight', 'f4')]),
+         "str320,int32,float32"),
+        # Jax
+        (jnp.array([[1, 2], [3, 4]]), "int32"),
+        (jnp.array([[1, 2], [3, 4]], dtype="int8"), "int8"),
+        # Tensorflow
+        (tf.constant([[1, 2], [3, 4]]), "int32"),
+        (tf.constant([[1, 2], [3, 4]], dtype="int64"), "int64"),
+        # Pytorch
+        (torch.tensor([[1, 2], [3, 4]]), "int64"),
+        (torch.tensor([[1, 2], [3, 4]], dtype=torch.int32), "int32"),
+    ],
+)
+def test_dtypes(value, expected):
+    assert ts.analysis._dtype(value) == expected

testing/test_dtype.py

@@ -48,7 +48,8 @@ def __init__(self,
                  vectorcolor="#fefecd", char_sep_scale=1.8, fontcolor='#444443',
                  underline_color='#C2C2C2', ignored_color='#B4B4B4', error_op_color='#A40227',
                  show:(None,'viz')='viz',
-                 hush_errors=True):
+                 hush_errors=True,
+                 dtype_colors=None, dtype_precisions=None, dtype_alpha_range=None):
         """
         Augment tensor-related exceptions generated from numpy, pytorch, and tensorflow.
         Also display a visual representation of the offending Python line that
@@ -109,15 +110,23 @@ def __init__(self,
                             unhandled code caught by my parser are ignored. Turn this off
                             to see what the error messages are coming from my parser.
         :param show: Show visualization upon tensor error if show='viz'.
+        :param dtype_colors: map from dtype w/o precision like 'int' to color
+        :param dtype_precisions: list of bit precisions to colorize, such as [32,64,128]
+        :param dtype_alpha_range: all tensors of the same type are drawn to the same color,
+                                  and the alpha channel is used to show precision; the
+                                  smaller the bit size, the lower the alpha channel. You
+                                  can play with the range to get better visual dynamic range
+                                  depending on how many precisions you want to display.
         """
-        self.show = show
-        self.fontname, self.fontsize, self.dimfontname, self.dimfontsize, \
+        self.show, self.fontname, self.fontsize, self.dimfontname, self.dimfontsize, \
         self.matrixcolor, self.vectorcolor, self.char_sep_scale,\
         self.fontcolor, self.underline_color, self.ignored_color, \
-        self.error_op_color, self.hush_errors = \
-            fontname, fontsize, dimfontname, dimfontsize, \
+        self.error_op_color, self.hush_errors, \
+        self.dtype_colors, self.dtype_precisions, self.dtype_alpha_range = \
+            show, fontname, fontsize, dimfontname, dimfontsize, \
             matrixcolor, vectorcolor, char_sep_scale, \
-            fontcolor, underline_color, ignored_color, error_op_color, hush_errors
+            fontcolor, underline_color, ignored_color, error_op_color, hush_errors, \
+            dtype_colors, dtype_precisions, dtype_alpha_range
 
     def __enter__(self):
         self.frame = sys._getframe().f_back # where do we start tracking? Hmm...not sure we use this
@@ -146,7 +155,10 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
                                               self.char_sep_scale, self.fontcolor,
                                               self.underline_color, self.ignored_color,
                                               self.error_op_color,
-                                              hush_errors=self.hush_errors)
+                                              hush_errors=self.hush_errors,
+                                              dtype_colors=self.dtype_colors,
+                                              dtype_precisions=self.dtype_precisions,
+                                              dtype_alpha_range=self.dtype_alpha_range)
                 if self.view is not None: # Ignore if we can't process code causing exception (I use a subparser)
                     if self.show=='viz':
                         self.view.show()
@@ -159,8 +171,8 @@ def __init__(self,
                  dimfontname='Arial', dimfontsize=9, matrixcolor="#cfe2d4",
                  vectorcolor="#fefecd", char_sep_scale=1.8, fontcolor='#444443',
                  underline_color='#C2C2C2', ignored_color='#B4B4B4', error_op_color='#A40227',
-                 savefig=None,
-                 hush_errors=True):
+                 savefig=None, hush_errors=True,
+                 dtype_colors=None, dtype_precisions=None, dtype_alpha_range=None):
         """
         As the Python virtual machine executes lines of code, generate a
         visualization for tensor-related expressions using from numpy, pytorch,
@@ -229,15 +241,23 @@ def __init__(self,
                             to see what the error messages are coming from my parser.
         :param savefig: A string indicating where to save the visualization; don't save
                         a file if None.
+        :param dtype_colors: map from dtype w/o precision like 'int' to color
+        :param dtype_precisions: list of bit precisions to colorize, such as [32,64,128]
+        :param dtype_alpha_range: all tensors of the same type are drawn to the same color,
+                                  and the alpha channel is used to show precision; the
+                                  smaller the bit size, the lower the alpha channel. You
+                                  can play with the range to get better visual dynamic range
+                                  depending on how many precisions you want to display.
         """
-        self.savefig = savefig
-        self.fontname, self.fontsize, self.dimfontname, self.dimfontsize, \
+        self.savefig, self.fontname, self.fontsize, self.dimfontname, self.dimfontsize, \
         self.matrixcolor, self.vectorcolor, self.char_sep_scale,\
         self.fontcolor, self.underline_color, self.ignored_color, \
-        self.error_op_color, self.hush_errors = \
-            fontname, fontsize, dimfontname, dimfontsize, \
+        self.error_op_color, self.hush_errors, \
+        self.dtype_colors, self.dtype_precisions, self.dtype_alpha_range = \
+            savefig, fontname, fontsize, dimfontname, dimfontsize, \
             matrixcolor, vectorcolor, char_sep_scale, \
-            fontcolor, underline_color, ignored_color, error_op_color, hush_errors
+            fontcolor, underline_color, ignored_color, error_op_color, hush_errors, \
+            dtype_colors, dtype_precisions, dtype_alpha_range
 
     def __enter__(self):
         # print("ON trace", sys._getframe())
@@ -301,6 +321,7 @@ def listener(self, frame, event, arg):
         filename, line = info.filename, info.lineno
         name = info.function
 
+        # Note: always true since L292 above...
         if event=='line':
             self.line_listener(module, name, filename, line, info, frame)
 
@@ -338,7 +359,10 @@ def viz_statement(self, code, frame):
                                  self.explainer.char_sep_scale, self.explainer.fontcolor,
                                  self.explainer.underline_color, self.explainer.ignored_color,
                                  self.explainer.error_op_color,
-                                 hush_errors=self.explainer.hush_errors)
+                                 hush_errors=self.explainer.hush_errors,
+                                 dtype_colors=self.explainer.dtype_colors,
+                                 dtype_precisions=self.explainer.dtype_precisions,
+                                 dtype_alpha_range=self.explainer.dtype_alpha_range)
         self.views.append(view)
         if self.explainer.savefig is not None:
             file_path = Path(self.explainer.savefig)
@@ -495,6 +519,23 @@ def istensor(x):
     return _shape(x) is not None
 
 
+def _dtype(v) -> str:
+    if hasattr(v, "dtype"):
+        dtype = v.dtype
+    elif "dtype" in v.__class__.__name__:
+        dtype = v
+    else:
+        return None
+
+    if dtype.__class__.__module__ == "torch":
+        # ugly but works
+        return str(dtype).replace("torch.", "")
+    if hasattr(dtype, "names") and dtype.names is not None and hasattr(dtype, "fields"):
+        # structured dtype: https://numpy.org/devdocs/user/basics.rec.html
+        return ",".join([_dtype(val) for val, _ in dtype.fields.values()])
+    return dtype.name
+
+
 def _shape(v):
     # do we have a shape and it answers len()? Should get stuff right.
     if hasattr(v, "shape") and hasattr(v.shape, "__len__"):

tsensor/analysis.py

@@ -21,4 +21,4 @@
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 """
-__version__ = '0.1.3'
+__version__ = '0.2'