    def _infer_shapes(self, x: Tensor) -> List[IntVar]:
        """Infers shapes for reduce ops."""

        input_dims = x._attrs["shape"]
        reduction_axes = self._attrs["reduction_axes"]
        if self._attrs["keepdim"]:
            output_dims = [
                IntImm(1) if idx in set(reduction_axes) else d
                for idx, d in enumerate(input_dims)
            ]
        else:
            # out codegen for reduce ops doesn't rely on the output shape,
            # so it's safe to squeeze the output tensor shape here
            output_dims = [
                d for idx, d in enumerate(input_dims) if idx not in set(reduction_axes)
            ]
        return output_dims
Infers shapes for reduce ops.
python/aitemplate/compiler/ops/reduce/reduce_common.py