Open brandonwillard opened 3 years ago
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
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.
Oops I was printing repeated entries from different modules. Updated the list (also the DeepCopy is already done in main)
Also, I don't think Scalar
Op
s should have infer_shape
methods.
I tried to filter them with and not isinstance(value, aesara.scalar.basic.ScalarOp)
It might help to have an explicit interface (e.g. a type/class) for Op
s that return TensorVariable
s—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.
_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?
The others I am not familiar with:
make_slice
,cpu_contiguous
, andlstsq
. 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.
We need to do a survey of the
Op
s and see which ones are still missingshape_infer
implementations. In general, we want to make sure that we're not unnecessarily evaluating graphs just to get their shapes.