search

tensorly / tensorly

TensorLy: Tensor Learning in Python.
http://tensorly.org
Other
1.51k stars 281 forks source link

Provide faster implementation of the MTTKRP #549

Closed aarmey closed 1 month ago

aarmey commented 1 month ago

Copied from #542, with just some minor word smithing.

cohenjer commented 1 month ago

Thanks @aarmey ! This looks perfect to me. "Only" need to sort this error with the entropy test...

aarmey commented 1 month ago

Ok, thanks @cohenjer. I'll mark it as broken and then proceed with this PR.

JeanKossaifi commented 1 month ago

@aarmey the reason we had changed this is that the faster version is not sparse safe - I guess the sparse backend needs a fair amount of work anyway.

cohenjer commented 1 month ago

@JeanKossaifi I was well aware of that that. That is why both methods are still in the package, and we allow the user to change the method easily with the backend manager. The procedure is in the documentation of both versions of mttkrp but I paste it here for further discussions: (here to switch from the khatri-rao version to the old version which I named unfolding_dot_khatri_rao_memory)

from tensorly.tenalg.core_tenalg.mttkrp import unfolding_dot_khatri_rao_memory
tl.tenalg.register_backend_method("unfolding_dot_khatri_rao", unfolding_dot_khatri_rao_memory)
tl.tenalg.use_dynamic_dispatch()

Using the backend manager, an expert user who requires the sparse backend (I guess this concerns very few people) can easily configure the mttkrp in his script to be the memory-efficient/sparse-safe/slow method. We can work on improving the sparse backend while enjoying fast dense computations by default. If you want we can swap back the default mttkrp, but in my opinion dense + fast should be the default one.

aarmey commented 1 month ago

The sparse backend could make this switch for the user, too.

Do the sparse backend tests not run? It would be helpful to codify which code is expected to be sparse safe that way.