    def _infer_shapes(self, x: Tensor) -> List[IntVar]:
        """Infer output shape for the softmax op.

        Parameters
        ----------
        x : Tensor
            Input shape of softmax op.

        Returns
        ----------
            List[IntVar]
        """
        shapes = x._attrs["shape"]
        assert (
            len(shapes) >= 2
        ), f"softmax only supports input with rank >= 2, current rank: {len(shapes)}"
        return x._attrs["shape"]
Infer output shape for the softmax op.

Parameters
----------
x : Tensor
Input shape of softmax op.

Returns
----------
List[IntVar]
python/aitemplate/compiler/ops/softmax/softmax.py