search

fferflo / einx

Universal Tensor Operations in Einstein-Inspired Notation for Python.
https://einx.readthedocs.io/en/stable/
MIT License
311 stars 8 forks source link

nested expression within `get_at` #1

Closed lucidrains closed 7 months ago

lucidrains commented 7 months ago

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 7 months ago

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 commented 7 months ago

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

fferflo commented 7 months ago

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