fferflo/einx

nested expression within `get_at`

lucidrains opened this issue · 3 comments

lucidrains commented

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 commented

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 4273732. Can you check if this works for you?

  • 🚀1
lucidrains commented

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

fferflo commented

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

  • 🎉1