search

april-tools / cirkit

a python framework to build, learn and reason about probabilistic circuits and tensor networks
https://cirkit-docs.readthedocs.io/en/latest/
GNU General Public License v3.0
71 stars 1 forks source link

Swap axes order -- data flow `(C)FKB` #103

Closed lkct closed 1 year ago

lkct commented 1 year ago

Based on #91, we should switch to a more efficient axes order.

The inner data flow is now (C)FKB, while the input and output of the model are BK which may be more intuitive. (output is contiguous as KB and a BK view is returned)

The input layer is not yet modified due to #93.

lkct commented 1 year ago

Benchmark results: 0.106/5.0

On the other hand, (C)FBK is 0.118/5.0, and therefore not chosen. (PS: with CFBK input, the einsum in mixing layer outputs non-contiguous tensor with all params configuration)

arranger1044 commented 1 year ago

This is a nice 2x speed up improvement! As a ref, in Antonio's experiments he was using collapsed CP, this might explain why he was (slightly) faster.

loreloc commented 1 year ago

Do we have benchmark results with a higher batch size? I feel that the choice of ...KB and ...BK depends on the batch size. I think ...BK is a more natural layout for our tensors. Moreover, this would remove the need of the final transpose on the circuit's output.

lkct commented 1 year ago

yes we should add more configurations to benchmark, and that's why I'm keeping a bak branch for FBK