tensor-compiler / taco

The Tensor Algebra Compiler (taco) computes sparse tensor expressions on CPUs and GPUs
http://tensor-compiler.org
Other
1.23k stars 184 forks source link

Adding dense tensors produce sparse tensors in pytaco #368

Open LinjianMa opened 3 years ago

LinjianMa commented 3 years ago

The following code

import pytaco as pt
a = pt.tensor([3,3], pt.dense)
b = pt.tensor([3,3], pt.dense)
c = a + b

produces c whose format is Format(({compressed},{compressed}; 0,1)) rather than dense.

Does it make sense? I think it's more reasonable to output dense tensors under this case?

rawnhenry commented 3 years ago

Yea, the output should definitely be dense. Figuring out what the output format should be was an issue we had when making APIs you are using for the taco python bindings. For now, it will just default to a sparse format in every mode of the tensor. I think you can change this by setting pt.default_mode = pt.dense.

For more control, you can use the API pt.tensor_add(a, b, [pt.dense, pt.dense]) to explicitly give the output format for a given operation. I would suggest just using that API instead for your use case. If the taco python bindings do the wrong thing here, that's a bug.

Note: I haven't tried running the commands above but I think they should work.

The reason taco doesn't do something more clever is just due to time constraints we had when implementing the bindings. I think one approach would be to reason about the output format level by level based on the operator being used. For example, for the plus operator, dense + dense = dense since addition is a union op. Something like sparse + sparse is a little unclear since you may want a dense result in some cases (eg. your inputs are disjoint and not very sparse) and a sparse result for others. For multiplication sparse * dense = sparse since its an intersection etc. Something like this would be easier to implement if each operator had an algebra baked into it (I have a PR for this coming soonTM). Maybe someone can revisit this when that PR is merged. Although, this approach has the caveat that it has to default to dense whenever op(0, 0) != 0 (due to the underlying assumption that tensors are sparse over 0).