aesara
aesara copied to clipboard
`local_shape_to_shape_i` removes `Assert`s
Here's a MWE:
import aesara
import aesara.tensor as at
from aesara.raise_op import assert_op
x = at.vector("x")
x_a = assert_op(x, at.gt(x.shape[0], 1))
y = x_a.shape[0]
aesara.dprint(y, print_type=True)
# Subtensor{int64} [id A] <TensorType(int64, ())>
# |Shape [id B] <TensorType(int64, (1,))>
# | |Assert{msg=Aesara Assert failed!} [id C] <TensorType(float64, (None,))>
# | |x [id D] <TensorType(float64, (None,))>
# | |Elemwise{gt,no_inplace} [id E] <TensorType(bool, ())>
# | |Subtensor{int64} [id F] <TensorType(int64, ())>
# | | |Shape [id G] <TensorType(int64, (1,))>
# | | | |x [id D] <TensorType(float64, (None,))>
# | | |ScalarConstant{0} [id H] <int64>
# | |TensorConstant{1} [id I] <TensorType(int8, ())>
# |ScalarConstant{0} [id J] <int64>
with aesara.config.change_flags(optimizer_verbose=True):
y_fn = aesara.function([x], y)
# optimizer: rewrite MergeOptimizer replaces ScalarConstant{0} of None with ScalarConstant{0} of None
# optimizer: rewrite local_subtensor_remove_broadcastable_index replaces Subtensor{int64}.0 of Subtensor{int64}(Shape.0, ScalarConstant{0}) with InplaceDimShuffle{}.0 of InplaceDimShuffle{}(Shape.0)
# optimizer: rewrite local_shape_to_shape_i replaces Shape.0 of Shape(Assert{msg=Aesara Assert failed!}.0) with MakeVector{dtype='int64'}.0 of MakeVector{dtype='int64'}(Shape_i{0}.0)
# optimizer: rewrite local_useless_dimshuffle_makevector replaces InplaceDimShuffle{}.0 of InplaceDimShuffle{}(MakeVector{dtype='int64'}.0) with Shape_i{0}.0 of Shape_i{0}(x)
# `local_shape_to_shape_i` appears to have removed the `Assert`
aesara.dprint(y_fn, print_type=True)
# Shape_i{0} [id A] <TensorType(int64, ())> 0
# |x [id B] <TensorType(float64, (None,))>