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

        x_shape = x._attrs["shape"]
        rank = len(x_shape)
        if rank <= 0:
            raise RuntimeError("expected a non-scalar tensor")
        if dim >= rank:
            raise RuntimeError(
                f"split dim ({dim}) expected to be less than rank ({rank})"
            )
        num_splits = len(split_sizes)
        if num_splits < 1:
            raise RuntimeError(
                f"the number of splits expected >=0 but got {num_splits}"
            )
        split_dim_size = x_shape[dim]._attrs["values"][0]
        if sum(split_sizes) != split_dim_size:
            raise RuntimeError(
                f"sum of split_sizes ({split_sizes}) does not match split_dim_size ({split_dim_size})"
            )

        output_shapes = []
        for split_size in split_sizes:
            output_shape = x_shape[:dim] + [IntImm(split_size)] + x_shape[dim + 1 :]
            output_shapes.append(output_shape)

        return output_shapes
Infers shapes for split.
python/aitemplate/compiler/ops/tensor/split.py