Generalize aesara.tensor.linalg.cholesky beyond 2D arrays
#1012
Conversation
|
How does the numpy implementation perform compared to scipy (including the lower version that is being removed)? If it's vastly different we could keep using it for 2D cases. For higher cases is numpy faster than a Python loop calling the scipy function? |
|
Besides performance concerns, the other thing you need to check is if there are any rewrites in Aesara that rely on the assumption that this Op is always 2-dimensional. If that's the case you might need to add a new check in the rewrites or generalize them to be higher-dimensional if that's already possible (and open a Aesara issue if not) |
| @@ -451,23 +421,6 @@ def test_solve_dtype(self): | |||
| assert x.dtype == x_result.dtype | |||
|
|
|||
|
|
|||
| def test_cho_solve(): | |||
There was a problem hiding this comment.
Why is this test removed?
|
|
||
|
|
||
| cho_solve = CholeskySolve() | ||
|
|
||
|
|
||
| def cho_solve(c_and_lower, b, check_finite=True): | ||
| def cho_solve(a, b): |
There was a problem hiding this comment.
This needs some care, as it will simply break people's code.
| check_finite=self.check_finite, | ||
| ) | ||
| out_dtype = aes.upcast(a.dtype, b.dtype) | ||
| broadcastable = [ |
There was a problem hiding this comment.
We can do better than broadcastable. Have a look at https://aesara.readthedocs.io/en/latest/extending/type.html
| @@ -39,10 +40,9 @@ class Cholesky(Op): | |||
| # TODO: for specific dtypes | |||
| # TODO: LAPACK wrapper with in-place behavior, for solve also | |||
|
|
|||
| __props__ = ("lower", "destructive", "on_error") | |||
| __props__ = ("destructive", "on_error") | |||
There was a problem hiding this comment.
Do we need destructive attribute here?
According to the Op docs, __props__ lists the attributes which influence the computation performed. And I am not sure how destructive attribute is used in perform method.
There was a problem hiding this comment.
My guess is that this was added for some functionality that was never implemented (e.g., a faster method that alters the input variables in place). Looking around the library I don't think this is being used anywhere so maybe we can remove it.
|
By the way, the rewrites involving linalg Ops are found here: https://github.com/aesara-devs/aesara/blob/main/aesara/sandbox/linalg/ops.py We should probably bring them over to this module (CC @brandonwillard). Importantly, some rewrites may need to be updated with the new Op's signatures... and the same performance concerns should be at least briefly investigated. |
brandonwillard
changed the title
aesara.tensor.linalg.cholesky beyond 2D arrays
brandonwillard
added
enhancement
See #499. |
Thank you for opening a PR!
Here are a few important guidelines and requirements to check before your PR can be merged:
pre-commitis installed and set up.Don't worry, your PR doesn't need to be in perfect order to submit it. As development progresses and/or reviewers request changes, you can always rewrite the history of your feature/PR branches.
If your PR is an ongoing effort and you would like to involve us in the process, simply make it a draft PR.