Add stablehlo.einsum, stablehlo.unary_einsum diffrules

[mofeing]

Fixes #82

@wsmoses I don't know much about LLVM, so I don't know how to implement them. But I wrote the prototypes and some annotations on how to continue the implementation.

[wsmoses]

You should copy the gradpad test which computes derivatives

[wsmoses]

@mofeing look here for example of how to deal with the einsum config.

Enzyme-JAX/src/enzyme_ad/jax/Implementations/HLODerivatives.td

Line 362 in b357c8d

def InversePermutation : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, [{

I can help as well, but my question for you is offhand do you know the corresponding derivative rule for the config

[mofeing]

I can help as well, but my question for you is offhand do you know the corresponding derivative rule for the config

@wsmoses The reverse derivative for einsum is very easy. Basically it's just another einsum per input where you replace the computing . Imagine the following einsum:

$$ C = \mathtt{einsum}(''$ia,$ib->$ic'', A, B) $$

where $A,B,C$ are tensors and $ia,ib,ic$ denote the list of indices of each tensor. The string is the einsum config.

Then the reverse derivatives are the following:

$$ dA = \mathtt{einsum}(''$ib,$ic->$ia'', conj(B), dC) $$

$$ dB = \mathtt{einsum}(''$ia,$ic->$ib'', conj(A), dC) $$

I see that we need to do the following:

1. Modify the einsum config like the patterns above.
2. Apply the conjugate only if A or B respectively, are complex. The main reason is that complex.conj only works on complex types, not on f32/f64.

[wsmoses]

To start with let's not add complex stuff yet and do that in a subsequent PR. Other things will also need to be updated here as well so might as well do it generically for everything.

[wsmoses]

[whoops meant to comment, not close...my b]

[wsmoses]

In any case, want to try to add the config update?

[mofeing]

Yeah, I see that GlobalExpr lets you write C++ code, right? I just don't know what I have access to and the methods of op inside that code.

For example, where is .getPermutation() defined in the code you pointed?

[wsmoses]

Yeah. You have op which is the current operation so it’s essentially the MLIR c++ api for your op

…

On Tue, May 28, 2024 at 3:17 PM Sergio Sánchez Ramírez < ***@***.***> wrote: Yeah, I see that GlobalExpr lets you write C++ code, right? I just don't know what I have access to and the methods of op inside that code. — Reply to this email directly, view it on GitHub <#83 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJTUXHYXP3QPXO3BYPOJYTZER7VLAVCNFSM6AAAAABIK6Q4HCVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCMZVGE4TKNJRGE> . You are receiving this because you modified the open/close state.Message ID: ***@***.***>

[mofeing]

@wsmoses just uploaded the stablehlo.einsum test with complex data. it might still error due to the result not being exactly the same, but it should be easily fixable.

To do

•  add activity support to stablehlo.einsum and stablehlo.unary_einsum diff rules
• fix diff rule for stablehlo.einsum on complex data

[wsmoses]

@mofeing tests seem to fail, worth investigating?