.getOutput());

    SmallVector<Type, 4> block_shape_splits_types(
        block_rank,
        tensorflow::GetTypeFromTFTensorShape({1}, rewriter.getIntegerType(64)));
    SmallVector<Value, 4> block_shape_splits(
        rewriter
            .create<SplitOp>(loc, block_shape_splits_types, zero_i32,
                             block_shape_i64)
            .getOutput());