    def _infer_shapes(self, x: Tensor):
        """Infers shapes for pad_last_dim."""

        x_shape = x._attrs["shape"]
        ndim = len(x_shape)
        if self._attrs["out_dim"] <= max(x_shape[-1]._attrs["values"]):
            raise RuntimeError("Output of padded dim must be larger than original dim")
        if ndim != self._attrs["ndim"]:
            raise RuntimeError("Data/Op dims mismatch")
        if ndim > 4:
            raise NotImplementedError
        output_shape = list(x_shape)
        output_shape[-1] = IntImm(self._attrs["out_dim"])
        return output_shape
Infers shapes for pad_last_dim.
python/aitemplate/compiler/ops/padding/pad_last_dim.py