Add einsum

[jessegrabowski]

Description

TODO:

• Support ellipsis (related to contract_path with ellipsis fails when shapes=True dgasmith/opt_einsum#235)
• Exclude broadcastable dims for better perf (JAX does the same)
• Handle missing static shape information (default to left to right contraction?)
  • Add rewrite for optimizing Einsum Ops when all inputs have known static shapes
  • Add rewrite for inlining optimized Einsum
• Get rid of Blockwise Reshape
  • Fix lingering infinite rewriting bug
• Decide on providing optimize kwarg
• Appease Mypy
• Better docstrings (@jessegrabowski self-assigned)
• Fix failing tests (@ricardoV94 self-assigned)

Related Issue

• Closes Implement einsum equivalent #57

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[zaxtax]

Are the current tests failing suppose to fail?

[jessegrabowski]

Looks like it's related to the changes to the __len__ method in variables.py. I'd suggest just reverting the change unless we really need it. It's a pretty low-level thing that would need a bit of work to figure out everything implicated.

[codecov]

Codecov Report

Attention: Patch coverage is 96.58793% with 13 lines in your changes missing coverage. Please review.

Project coverage is 81.69%. Comparing base (58fec45) to head (de06ff6).
Report is 99 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/einsum.py	96.90%	3 Missing and 3 partials ⚠️
pytensor/link/jax/dispatch/einsum.py	55.55%	4 Missing ⚠️
pytensor/tensor/rewriting/basic.py	96.61%	1 Missing and 1 partial ⚠️
pytensor/tensor/rewriting/shape.py	96.15%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #722      +/-   ##
==========================================
+ Coverage   81.60%   81.69%   +0.09%     
==========================================
  Files         179      182       +3     
  Lines       47271    47585     +314     
  Branches    11481    11584     +103     
==========================================
+ Hits        38574    38875     +301     
- Misses       6511     6520       +9     
- Partials     2186     2190       +4     

Files with missing lines	Coverage Δ
pytensor/link/jax/dispatch/__init__.py	100.00% <100.00%> (ø)
pytensor/tensor/basic.py	91.41% <100.00%> (+0.02%)	⬆️
pytensor/tensor/functional.py	100.00% <100.00%> (ø)
pytensor/tensor/rewriting/blockwise.py	97.32% <100.00%> (+2.52%)	⬆️
pytensor/tensor/rewriting/einsum.py	100.00% <100.00%> (ø)
pytensor/tensor/rewriting/ofg.py	100.00% <100.00%> (ø)
pytensor/tensor/shape.py	92.69% <100.00%> (+0.01%)	⬆️
pytensor/tensor/rewriting/shape.py	82.74% <96.15%> (+0.71%)	⬆️
pytensor/tensor/rewriting/basic.py	94.15% <96.61%> (+0.25%)	⬆️
pytensor/link/jax/dispatch/einsum.py	55.55% <55.55%> (ø)
... and 1 more

... and 2 files with indirect coverage changes

[ricardoV94]

All cases except those requiring tensordot with batch dims not on the left are passing

We may need more tests soon enough

[zaxtax]

Can we reuse how numpy implements tensordot?

…

On Tue, 7 May 2024, 11:52 Ricardo Vieira, ***@***.***> wrote: All cases except those requiring tensordot with batch dims not on the left are passing We may need more tests soon enough — Reply to this email directly, view it on GitHub <#722 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUL3W5UGQ73SNQMCO2LZBEOQDAVCNFSM6AAAAABGPQHTOGVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDAOJZGA4DMOBVHE> . You are receiving this because your review was requested.Message ID: ***@***.***>

[ricardoV94]

Can we reuse how numpy implements tensordot?

We already do that, but numpy doesn't have batched tensordot (except of course through einsum), but we already have batched tensordot working in this PR just not with arbitrary batch axis. Should just need some extra transposes to get the job done

[ricardoV94]

Huhu convolutions via einsum work :D

[ricardoV94]

Now Einsum also works with inputs with unknown static shape (unoptimized ofc). We can add a rewrite for when such Op is found with inputs that now have static shapes (this can be quite relevant in PyMC, when users use freeze_rv_and_dims on a model with mutable coords)

[ricardoV94]

The ellipsis case is failing due to a bug in opt_einsum: dgasmith/opt_einsum#235

[zaxtax]

Oh I misunderstood!

…

On Thu, 11 Jul 2024, 12:38 Ricardo Vieira, ***@***.***> wrote: Not sure what you mean @zaxtax <https://github.com/zaxtax>, I'm talking about allowing the "optimize" kwarg like there is in numpy, which defines what kind of optimization to do: optimize{bool, list, tuple, ‘greedy’, ‘optimal’}, users can pass their custom contraction path as well. If users pass contraction_path, we don't need to know static shapes. If users set to greedy/optimal (optimal should be default), we need to know. But we may find them later only. If they don't want optimize, then we don't need to obviously — Reply to this email directly, view it on GitHub <#722 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUPHF22H4SECYZCOLV3ZLZOBFAVCNFSM6AAAAABGPQHTOGVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDEMRSGU4TIOJZHE> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[ricardoV94]

Some unrelated jax test failing, probably something that changed in a recent release? https://github.com/pymc-devs/pytensor/actions/runs/10161294793/job/28099514375?pr=722#step:6:778

[jessegrabowski]

Yes, I am also looking at this now. It's a jax bug that can be recreated easily:

import jax
jax.jit(jax.numpy.tri)(3, 3, 0)

We can ignore it. Looks like their _canonicalize_axis function is underflowing (something like np.uint32(-1) )

[ricardoV94]

I think I fixed the tests (not the JAX one) and appeased mypy. @jessegrabowski docstrings extensions are left to you

[ricardoV94]

Stopped force-pushing if you want to take over

[jessegrabowski]

Opened an issue here: jax-ml/jax#22751

I'll hit the docstrings ASAP if that's all that's holding this up

[ricardoV94]

Opened an issue here: jax-ml/jax#22751

Great let's just mark it as xfail then

[jessegrabowski]

First pass on docstrings. Working on the doctests revealed two things:

1. Our implementation of _delta does not agree with that of JAX:

from jax._src.lax.lax import _delta as jax_delta
from pytensor.tensor.einsum import _delta as pt_delta
jax_delta(int, (3, 3, 3), (0,1))

Array([[[1, 1, 1],
        [0, 0, 0],
        [0, 0, 0]],

       [[0, 0, 0],
        [1, 1, 1],
        [0, 0, 0]],

       [[0, 0, 0],
        [0, 0, 0],
        [1, 1, 1]]], dtype=int32)


pt_delta((3,3,3), (0,1)).astype(int).eval()
array([[[1, 0, 0],
        [0, 1, 0],
        [0, 0, 1]],

       [[1, 0, 0],
        [0, 1, 0],
        [0, 0, 1]],

       [[1, 0, 0],
        [0, 1, 0],
        [0, 0, 1]]])

We seem to always output the axes=(1,2) case, regardless of what the requested axes were.

2. Our _general_dot function is discarding shape information somewhere, which doesn't seem right:

        import pytensor.tensor as pt
        from pytensor.tensor.einsum import _general_dot
        A = pt.tensor(shape = (3, 4, 5))
        B = pt.tensor(shape = (3, 5, 2))

        result = _general_dot((A, B), axes=[[2], [1]], batch_axes=[[0], [0]])
        print(result.type.shape)

       (3, None, None)

[ricardoV94]

Static shape is optional, not a requirement. In our case it probably has to do with the reshape introduced by tensordot and/or Blockwise which doesn't do any special shape inference shape (static or at rewrite) for core shapes.

That's something we probably want to address for Blockwise in the Numba backend

[jessegrabowski]

I understand it's optional, but it also shouldn't be discarded if available no?

[ricardoV94]

We are not discarding anything on purpose but an intermediate op (or blockwise) doesn't know how to provide more precise output shape.

There can also be a tradeoff where quite some effort may be needed to figure out static shape that may not be worth it at define time. Anyway the main point is that it shouldn't be a blocker.

We can open an issue for whatever Op is losing the static shape and then assess if it's worth the cost or not

[ricardoV94]

@jessegrabowski I think I fixed the _delta, it was missing the equivalent call to broadcast_in_dim

[zaxtax]

🎆