def test_fuse_squeeze_expand_dims_connected_output(self):
        """
        If squeeze is connected to block output, it cannot be removed.
        However, the expand_dims can be a block output.
        """
        # squeeze connected to output. Nothing happens.
        @mb.program(input_specs=[mb.TensorSpec(shape=(1,))])
        def prog(x):
            squeeze = mb.squeeze(x=x, axes=(0,))
            expand_dims = mb.expand_dims(x=squeeze, axes=(0,))
            return mb.relu(x=expand_dims), squeeze

        apply_pass_and_basic_check(prog, "common::fuse_squeeze_expand_dims")
        assert get_op_types_in_program(prog) == ["squeeze", "expand_dims", "relu"]

        # expand_dims connected to output. Still good to fuse.
        @mb.program(input_specs=[mb.TensorSpec(shape=(1,))])
        def prog(x):
            squeeze = mb.squeeze(x=x, axes=(0,))
            expand_dims = mb.expand_dims(x=squeeze, axes=(0,))
            return mb.relu(x=expand_dims), expand_dims

        apply_pass_and_basic_check(prog, "common::fuse_squeeze_expand_dims")
        assert get_op_types_in_program(prog) == ["identity", "relu"]

        If squeeze is connected to block output, it cannot be removed.
        However, the expand_dims can be a block output.
        
coremltools/converters/mil/mil/passes/tests/test_passes.py