parrt / tensor-sensor

The goal of this library is to generate more helpful exception messages for matrix algebra expressions for numpy, pytorch, jax, tensorflow, keras, fastai.
https://github.com/parrt/tensor-sensor
MIT License
771 stars 39 forks source link

Showing too many matrices for complicated operands #32

Open parrt opened 3 years ago

parrt commented 3 years ago

The following code generates an exception but instead of showing the result of the operand subexpressions, it shows all bits of it:

toomany-1

import torch
import tsensor

n = 200         # number of instances
d = 764         # number of instance features
nhidden = 256

Whh = torch.eye(nhidden, nhidden)  # Identity matrix
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
r = torch.randn(nhidden, 3)         # fake this computation
X = torch.rand(n,d)                 # fake input

# Following code raises an exception
with tsensor.clarify():
    h = torch.tanh(Whh @ (r*h) + Uxh @ X.T + bh)  # state vector update equation