def _infer_shapes(self, x: Tensor, w: Tensor):
        """infer the output shape for attention module"""
        x_shape_values = [var._attrs["values"] for var in x._attrs["shape"]]
        x_shapes = itertools.product(*x_shape_values)
        w_shape = [var._attrs["values"][0] for var in w._attrs["shape"]]
        # run infer shape for each
        y_shapes = []
        for x_shape in x_shapes:
            y_shape = self._infer_shape(x_shape, w_shape)
            y_shapes.append(y_shape)

        def unique(vector):
            return sorted(set(vector))

        batch_info = x._attrs["shape"][0]
        output_shape = [
            batch_info,
            x._attrs["shape"][2],
            x._attrs["shape"][1],
            w._attrs["shape"][-1],
        ]
        return output_shape
infer the output shape for attention module
python/aitemplate/compiler/ops/attention/mem_eff_attention.py