aesara icon indicating copy to clipboard operation
aesara copied to clipboard

`local_shape_to_shape_i` removes `Assert`s

Open brandonwillard opened this issue 1 year ago • 0 comments

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,))>

brandonwillard avatar Aug 09 '22 00:08 brandonwillard