Track down `score_matching.py` slowdown issue

gchq / coreax

A library for coreset algorithms, written in Jax for fast execution and GPU support.

Apache License 2.0

25 stars 2 forks source link

Track down `score_matching.py` slowdown issue #764

Closed db091756 closed 1 month ago

db091756 commented 2 months ago

What's the new feature?

Changing how enforce 2-dimensionality in the score_matching.py code slows down SteinThinning very significantly. Track down this bug. (See https://github.com/gchq/coreax/pull/746#issuecomment-2321449604)

What value does this add?

Consistent dimensionality handling without slowdown!

Is there an alternative you've considered?

No response

Additional context

No response

tc85324 commented 1 month ago

Was this resolved when #746 was merged?

db091756 commented 1 month ago

Was this resolved when #746 was merged?

We reverted the line of code that caused the issue so it is as fast as it was previously. But it is still worth understanding where the issue is. I believe it is something to do with computing the hessian but not looked into it fully