Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adding Subtensor static shape inference in Subtensor.make_node #935

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

larryshamalama
Copy link
Contributor

Closes #922.

As per @brandonwillard's comment, I renamed broadcastable to static_shape, incorporating the other suggested minor edits.

I took inspiration from lines 773 to 777 from Subtensor.infer_shape for shape inference.

No new tests have been added yet as I am creating this PR.

@larryshamalama
Copy link
Contributor Author

larryshamalama commented Apr 30, 2022

N.B.: Local testing yields CompilationErrors, probably (?) because mkl provides a C++ interface via Python, except installation of the package is difficult on an M1 Mac (c.f. issue #909). Would need to think of a (possibly hackish) solution, at least for myself

Copy link
Member

@brandonwillard brandonwillard left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The error in CI says that the rewrite local_subtensor_make_vector introduced a type error when it tried to replace a variable with another variable that has an incompatible static shape.

Inspection via the debugger shows that the node being rewritten has the following form:

Subtensor{int64::} [id BH] <TensorType(int64, (5,))> '' 4

| MakeVector{dtype='int64'} [id BD] <TensorType(int64, (3,))> '' 3 |
| ScalarConstant{-2} [id P] |

A shape of (5,) for the output of a Subtensor Op like x[-2::], where x is a vector of length 3, doesn't make much sense.

Let's see if we can reproduce the error by creating an equivalent graph.

import aesara
import aesara.tensor as at


# A vector of length three built from a `MakeVector` `Op`
x = at.as_tensor([at.lscalar(), at.lscalar(), at.lscalar()])
z = x[-2:]

aesara.dprint(z, print_type=True, depth=2)
# Subtensor{int64::} [id A] <TensorType(int64, (5,))> ''
#  |MakeVector{dtype='int64'} [id B] <TensorType(int64, (3,))> ''
#  |ScalarConstant{-2} [id C] <int64>

We've successfully reproduced the problematic graph, so we can now isolate the issue in Subtensor.make_node.

Also, we now have a good unit test to add to tests.tensor.test_subtensor. Since all the current tests in tests.tensor.test_subtensor seem to pass when this bug is present, it's clear that those tests are lacking coverage for some simple and important cases. One of which is covered by the simple example above.

FYI: Simple test additions like these can be just as valuable to Aesara as the feature work from which they are derived.

@larryshamalama
Copy link
Contributor Author

Thanks for the detailed input. I see where the error is and will try to fix it. Maybe I can reuse get_canonical_form_slice in some way, but I'm working on it.

Right now I am spending time addressing issues with testing on my Mac. I was running into graph compilation issues in pytest but it turns out that they are due to the lack of mkl...

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working shape inference
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Subtensor shape incorrect shape inference
3 participants