def extract(cls, node: Node):
        attrs = {
            'begin_mask': node.module.begin_mask,
            'end_mask': node.module.end_mask,
            'shrink_axis_mask': node.module.shrink_axis_mask,
            'new_axis_mask': [0],
            'ellipsis_mask': [0],
        }

        StridedSlice.update_node_stat(node, attrs)
        return cls.enabled
Example #2
0
    def extract(cls, node):
        pb = node.pb
        bm = int_to_array_bit_mask(pb.attr["begin_mask"].i)
        bm = np.array([1 - b for b in bm], dtype=np.int32)
        em = int_to_array_bit_mask(pb.attr["end_mask"].i)
        em = np.array([1 - b for b in em], dtype=np.int32)
        attrs = {
            'begin_mask':
            bm,
            'end_mask':
            em,
            'ellipsis_mask':
            int_to_array_bit_mask(pb.attr["ellipsis_mask"].i),
            'new_axis_mask':
            int_to_array_bit_mask(pb.attr["new_axis_mask"].i),
            'shrink_axis_mask':
            int_to_array_bit_mask(pb.attr["shrink_axis_mask"].i),
        }

        StridedSlice.update_node_stat(node, attrs)
        return cls.enabled