nested expression within `get_at`

[lucidrains]

hello Florian, and thank you for working on this project!

i've been testing out get_at and it is working like a charm! will probably start using it more and more just for this function alone

wanted to consult you about a special case that I thought einx should be able to handle, but errors out on. it could be due to my not knowing how to write the right expression

the piece of code i'm trying to replace is this with something like get_at('b t [d], b (t [1]) -> b (t [1])', prob, indices). this is a pattern that is recurring throughout a bunch of works, so would be nice if there were an einx.get_at one-liner!

[fferflo]

Hi Phil, thanks for your interest!

For 1D indexing one can simply leave out the coordinate axis in the indices expression, so this should work:

get_at("b t [d], b t -> b t", prob, indices)

or even

get_at("... [d], ... -> ...", prob, indices)

Although when testing this I found a bug that probably caused the same error you are also seeing. It should be fixed with https://github.com/fferflo/einx/commit/4273732b401a00f02dceedd80b65d5bb485eded0. Can you check if this works for you?

[lucidrains]

@fferflo lightspeed response! it works! thank you Florian; dinner or beer is on me if you are ever in SF 😃

[fferflo]

Thanks, glad it works and I'll keep it in mind for my next (first) trip to SF 😄