Batch Data Processing Module

A module for processing large batches of multi-dimensional data by dividing them into smaller, manageable chunks.

Overview

You need to implement a batch data processor that can split large arrays into smaller sub-batches for distributed processing. The module should handle 1D vectors, 2D matrices, and 3D tensors, providing flexible splitting strategies along different axes.

Requirements

Core Functionality

The module must provide the following splitting capabilities:

1. Basic splitting: Split an array into equal-sized chunks along a specified axis
2. Flexible splitting: Split an array into chunks of potentially unequal sizes when even division is not possible
3. Column-wise splitting: Split 2D data by separating columns
4. Row-wise splitting: Split 2D data by separating rows
5. Depth-wise splitting: Split 3D data along the third dimension

Input Specifications

• Accept GPU arrays as input for processing
• Support 1D, 2D, and 3D input arrays
• Handle cases where the array size is not evenly divisible by the number of chunks

Output Requirements

• Return a list of sub-arrays after splitting
• Ensure all sub-arrays are also GPU arrays
• Preserve data types and values during splitting

Test Cases

• Given a 1D array of 10 elements, splitting it into 3 chunks produces 3 sub-arrays with sizes [4, 3, 3] @test
• Given a 2D array with shape (6, 8), splitting horizontally into 4 chunks produces 4 sub-arrays each with shape (6, 2) @test
• Given a 2D array with shape (9, 4), splitting vertically into 3 chunks produces 3 sub-arrays each with shape (3, 4) @test
• Given a 3D array with shape (4, 5, 12), splitting along the depth dimension into 3 chunks produces 3 sub-arrays each with shape (4, 5, 4) @test

Implementation

@generates

API

def split_into_chunks(array, num_chunks, axis=0):
    """
    Split an array into the specified number of chunks along the given axis.
    Handles cases where array cannot be evenly divided.

    Parameters:
    - array: Input array to split
    - num_chunks: int, number of chunks to create
    - axis: int, axis along which to split (default: 0)

    Returns:
    list: List of sub-arrays
    """

def split_columns(array, num_chunks):
    """
    Split a 2D array into chunks by dividing columns.

    Parameters:
    - array: 2D input array
    - num_chunks: int, number of column chunks to create

    Returns:
    list: List of sub-arrays with divided columns
    """

def split_rows(array, num_chunks):
    """
    Split a 2D array into chunks by dividing rows.

    Parameters:
    - array: 2D input array
    - num_chunks: int, number of row chunks to create

    Returns:
    list: List of sub-arrays with divided rows
    """

def split_depth(array, num_chunks):
    """
    Split a 3D array into chunks along the depth (third) dimension.

    Parameters:
    - array: 3D input array
    - num_chunks: int, number of depth chunks to create

    Returns:
    list: List of sub-arrays with divided depth
    """

Dependencies { .dependencies }

cupy { .dependency }

Provides GPU-accelerated array operations and splitting functions.