Remove `block_diag` from `pymc.math` in favor of alias to `pytensor.tensor.slinalg.block_diag`

[jessegrabowski]

Description

block_diag is a core math function that should live in pytensor and be aliased in pm.math, rather than being directly defined there. pymc-devs/pytensor#576 moves the block_diag code, along with relevant support code (pm.math.ix, pm.math.largest_common_dtype, into pytensor. The PyMC function should be replaced with an alias to that new function once merged.

In addition, the pytensor function is going to more closely follow the scipy API, which means that the function signature will change from pm.math.block_diag(matrices) to pm.math.block_diag(*matrices). Scipy also has separate functions for the sparse and dense cases, which I think we should have as well (rather than a single function). This necessitates a depreciation warning and a temporary helper function to dispatch to either pytensor.tensor.slinalg.block_diag or pytensor.sparse.block_diag

[OmGhadge]

Hi @jessegrabowski,

I have a question regarding the suggested steps mentioned in the issue description. It mentions both creating an alias in pm.math and having a temporary helper function. Based on my understanding, the alias alone seems sufficient for providing compatibility.

Would you clarify whether the intention is to deprecate pm.math.block_diagonal(matrices) with a warning (with existing implementation and no alias), allowing users to still use it while creating a separate helper function with an alias to pytensor.tensor.slinalg.block_diag or pytensor.sparse.block_diag?

Could you please clarify ?

[jessegrabowski]

I think what would be best would be to keep pm.math.block_diagonal, but it should just dispatch to either pt.linalg.block_diag or pytensor.spare.block_diag, depending on the types of the matrices. Users shouldn't notice a difference.

All of the code/machinery in pm.math related to block_diag could then just be deleted

[ricardoV94]

@OmGhadge are you interested in working on this one?

[OmGhadge]

Hi @ricardoV94 ,

Thank you for reaching out. I genuinely appreciate your consideration. While I would have liked to work on resolving the issue, I see that @AryanNanda17 has already submitted a pull request for it. I believe it's best for him to proceed, especially since the issue seems mostly resolved, with only some test coverage remaining.

However, I would like to work on other issues. In fact, I've opened a draft pull request (#613) that addresses another issue (#323) in the pytensor repository. I would greatly appreciate it if you could take some time to review it.

[ricardoV94]

Thanks @OmGhadge! Your PR is on my backlog, I haven't forgot about it! Will try to get to it sometime soon