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

        output_shape = list(x._attrs["shape"])
        output_shape[-1] = IntImm(self._attrs["topK"])
        return output_shape
Infers shapes for topK.
python/aitemplate/compiler/ops/tensor/topk.py