search

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

Suppress visualisation of () as operator in tree #22

Closed sv2518 closed 3 years ago

sv2518 commented 3 years ago

Hello.

I am using tensor-sensor to visualise DAGs of the domain specific language Slate, which is used in a code generation framework for Finite element methods called Firedrake. Slate expresses linear algebra operations on tensors. I am using tensor-sensor to visualise the DAG before and after an optimisation pass. An example would be the following:

Before optimisation: triplemul_beforeopt.pdf After optimisation: tripleopt_afteropt.pdf

While the visualisation of the tree is correct in both cases, I would quite like to suppress the node for the brackets (i.e. for the SubExpr node) to avoid confusion about the amount of temporaries generated. Is there already a way of controlling this as a user and if not would there be interest in supporting it?

Best wishes and thanks in advance, Sophia

parrt commented 3 years ago

Hi Sophia, Interesting. Do you want to remove the entire bracket node or just the () part displayed in the node? I would have to tweak the code to avoid this I would say; no easy option for the user. If you're familiar with parsing, perhaps you could take a look and isolate the issue in the code?

sv2518 commented 3 years ago

I would like to remove the whole node, so tbc, in the picture after the optimisation I would only like to see three blue boxes. I had a first go at this before I opened the issue, but it didn't seem straightforward. I have just tried again and came up with a solution. I have only added some tweaking in the visualisation function and not in the parsing, so that there is no loss of information about the order of the information by dropping brackets. I have added support to drive this feature as a user, up to you if you would like to support this or not. I will open a PR.

sv2518 commented 3 years ago

Thanks!