GEM: Simplify Indexed tensors #131
Conversation
pbrubeck
force-pushed
the
pbrubeck/simplify-indexed
branch
from
4cff9b4 to
783e0c4
Compare
pbrubeck
force-pushed
the
pbrubeck/simplify-indexed
branch
from
33c66b4 to
f445d5b
Compare
pbrubeck
force-pushed
the
pbrubeck/simplify-indexed
branch
from
f445d5b to
de501c1
Compare
pbrubeck
force-pushed
the
pbrubeck/simplify-indexed
branch
from
d2e362b to
592061e
Compare
pbrubeck
force-pushed
the
pbrubeck/simplify-indexed
branch
from
1982373 to
4ab1c0d
Compare
pbrubeck
force-pushed
the
pbrubeck/simplify-indexed
branch
from
4ab1c0d to
b49eb5c
Compare
|
I marked this PR as Draft as you still seem to be debugging. Please mark it as Ready for review when you become confident. |
pbrubeck
force-pushed
the
pbrubeck/simplify-indexed
branch
from
996f2b4 to
5dfd17d
Compare
|
Tests are passing (modulo known doc/linkcheck failures) https://github.com/firedrakeproject/firedrake/actions/runs/17322729298/job/49186947155 @ksagiyam I feel quite happy with this, could you review? |
| if isinstance(B, Indexed): | ||
| C, = B.children | ||
| kk = B.multiindex | ||
| if not isinstance(C, ComponentTensor): |
There was a problem hiding this comment.
Why do we need to avoid ComponentTensor here?
There was a problem hiding this comment.
I think it is because for ComponentTensor we cannot replace indices if the substitution involves the free indices, for other classes this seems fine.
| if all(isinstance(elem, Indexed) for elem in array.flat): | ||
| tensor = e0.children[0] | ||
| multiindex = tuple(i for i in e0.multiindex if not isinstance(i, Integral)) | ||
| index_shape = tuple(i.extent for i in multiindex if isinstance(i, Index)) | ||
| if index_shape + array.shape + child_shape == tensor.shape: | ||
| if all(elem.children[0] == tensor for elem in array.flat[1:]): | ||
| if all(elem.multiindex == multiindex + idx for idx, elem in numpy.ndenumerate(array)): | ||
| return partial_indexed(tensor, multiindex) |
There was a problem hiding this comment.
Maybe we should make it more explicit that we are only handling multiindex of the following pattern: (Index, Index, Index, ..., Integral, Integral, ...). It looks the case where we have VariableIndexs, for instance, is handled in a very obscure way.
| slices = tuple(i if isinstance(i, int) else slice(None) for i in multiindex) | ||
| sub = aggregate.array[slices] | ||
| sub = Literal(sub, dtype=aggregate.dtype) if isinstance(aggregate, Constant) else ListTensor(sub) | ||
| return Indexed(sub, tuple(i for i in multiindex if not isinstance(i, int))) |
There was a problem hiding this comment.
I'm not sure if this is safe. This is a recursion, and, unlike when we use DAGTraverser, the result is not cached.
There was a problem hiding this comment.
What is the issue? Are we potentially creating new objects and not freeing memory for the old ones?
| if all((j in kk) and (j not in ff) for j in jj): | ||
| rep = dict(zip(jj, ii)) | ||
| ll = tuple(rep.get(k, k) for k in kk) | ||
| return Indexed(C, ll) |
No description provided.