Add missing `Op.infer_shape` implementations

[brandonwillard]

We need to do a survey of the Ops and see which ones are still missing shape_infer implementations. In general, we want to make sure that we're not unnecessarily evaluating graphs just to get their shapes.

• [x] aesara.compile.ops.DeepCopyOp` #560
• [ ] aesara.tensor.nlinalg.lstsq

[ricardoV94]

No idea how bad this is:

from types import ModuleType
import aesara

def find_ops_in_modules(module, modules_visited, ops):
    for key, value in module.__dict__.items():
        if isinstance(value, aesara.graph.op.Op):
            if (
                not hasattr(value, "infer_shape")
                and not isinstance(value, aesara.tensor.elemwise.Elemwise)
                and not isinstance(value, aesara.scalar.basic.ScalarOp)
                and key not in ops
            ):
                ops.add(key)
                print(module.__name__, key)
                continue
        if isinstance(value, ModuleType):
            if value not in modules_visited:
                modules_visited.add(value)
                find_ops_in_modules(value, modules_visited, ops)

ops = set()
modules_visited = set()
find_ops_in_modules(aesara, modules_visited, ops)

aesara.tensor.basic _nonzero
aesara.tensor.basic default
aesara.tensor.type_other make_slice
aesara.tensor.extra_ops cpu_contiguous
aesara.tensor.nlinalg lstsq

[brandonwillard]

That's exactly what we need to do; automate the search for things like this! It might be possible to use the Op MRO tree to get more coverage, but this is definitely the right approach.

Aside from lstsq (and possibly default), some of those are Variable types that don't have shapes (e.g. make_slice), as I recall.

[ricardoV94]

Oops I was printing repeated entries from different modules. Updated the list (also the DeepCopy is already done in main)

[brandonwillard]

Also, I don't think Scalar Ops should have infer_shape methods.

[ricardoV94]

I tried to filter them with and not isinstance(value, aesara.scalar.basic.ScalarOp)

[brandonwillard]

It might help to have an explicit interface (e.g. a type/class) for Ops that return TensorVariables—or at least things with a shape that can make use of infer_shape. Such an interface would need to offer/do more than just say that an Op could/should have an infer_shape, though.

[ricardoV94]

_nonzero and default shape depends on the actual op computation, so those don't make much sense.

The others I am not familiar with: make_slice, cpu_contiguous, and lstsq. Perhaps the last one?

[brandonwillard]

The others I am not familiar with: make_slice, cpu_contiguous, and lstsq. Perhaps the last one?

Yeah, that one should have an infer_shape. It's not a common/critical Op, so there's no priority for that one, but we should definitely put it on the list.