Showing too many matrices for complicated operands

[parrt]

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

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