generates balanced stratified synthetic test data to train a model on args trainsamples integer how many training samples to generate testsamples integer how many test samples to generate inputshape tuple of integers shape of the inputs numclasses integer number of classes for the data and targets randomseed integer random seed used by numpy to generate data returns a tuple of numpy arrays xtrain ytrain xtest ytest total samples ensure that we generate a balanced dataset generate extra samples in a deterministic manner generate data shuffle the entire dataset to ensure randomness based on seed stratified shuffle split convert to numpy arrays shuffle training and test sets after stratified split utility to generate the cartesian product of parameters values and generate a test case names for each combination the result of this function is to be used with the parameterized namedparameters decorator it is a replacement for parameterized product which adds explicit test case names for example this code class namedexampleparameterized testcase parameterized namedparameters namedproduct testcasename negative x 1 testcasename positive x 1 testcasename zero x 0 numeraltypefloat int def testconversionself x numeraltype self assertequalnumeraltypex x produces six tests note that absl will reorder them by name namedexample testconversionnegativefloat namedexample testconversionpositivefloat namedexample testconversionzerofloat namedexample testconversionnegativeint namedexample testconversionpositiveint namedexample testconversionzeroint this function is also useful in the case where there is no product to generate test case names for one argument parameterized namedparametersnamedproductnumeraltypefloat int args args each positional parameter is a sequence of keyword arg dicts every test case generated will include exactly one dict from each positional parameter these will then be merged to form an overall list of arguments for the test case each dict must contain a testcasename key whose value is combined with others to generate the test case name kwargs a mapping of parameter names and their possible values possible values should given as either a list or a tuple a string representation of each value is used to generate the test case name returns a list of maps for the test parameters combinations to pass to parameterized namedparameters convert the keyword arguments in the same dict format as the args the current list of tests start with one empty test augment the testcase name by appending augment the test by adding all the parameters overwrite the list of tests with the product obtained so far generates balanced stratified synthetic test data to train a model on args train_samples integer how many training samples to generate test_samples integer how many test samples to generate input_shape tuple of integers shape of the inputs num_classes integer number of classes for the data and targets random_seed integer random seed used by numpy to generate data returns a tuple of numpy arrays x_train y_train x_test y_test total samples ensure that we generate a balanced dataset generate extra samples in a deterministic manner generate data shuffle the entire dataset to ensure randomness based on seed stratified shuffle split convert to numpy arrays shuffle training and test sets after stratified split utility to generate the cartesian product of parameters values and generate a test case names for each combination the result of this function is to be used with the parameterized named_parameters decorator it is a replacement for parameterized product which adds explicit test case names for example this code class namedexample parameterized testcase parameterized named_parameters named_product testcase_name negative x 1 testcase_name positive x 1 testcase_name zero x 0 numeral_type float int def test_conversion self x numeral_type self assertequal numeral_type x x produces six tests note that absl will reorder them by name namedexample test_conversion_negative_float namedexample test_conversion_positive_float namedexample test_conversion_zero_float namedexample test_conversion_negative_int namedexample test_conversion_positive_int namedexample test_conversion_zero_int this function is also useful in the case where there is no product to generate test case names for one argument parameterized named_parameters named_product numeral_type float int args args each positional parameter is a sequence of keyword arg dicts every test case generated will include exactly one dict from each positional parameter these will then be merged to form an overall list of arguments for the test case each dict must contain a testcase_name key whose value is combined with others to generate the test case name kwargs a mapping of parameter names and their possible values possible values should given as either a list or a tuple a string representation of each value is used to generate the test case name returns a list of maps for the test parameters combinations to pass to parameterized named_parameters convert the keyword arguments in the same dict format as the args the current list of tests start with one empty test augment the testcase name by appending augment the test by adding all the parameters overwrite the list of tests with the product obtained so far
import numpy as np


def get_test_data(
    train_samples, test_samples, input_shape, num_classes, random_seed=None
):
    
    np.random.seed(random_seed)

    
    total_samples = train_samples + test_samples

    
    samples_per_class = total_samples // num_classes
    y = np.array(
        [i for i in range(num_classes) for _ in range(samples_per_class)],
        dtype=np.int32,
    )

    
    extra_samples = total_samples - len(y)
    y_extra = np.array(
        [i % num_classes for i in range(extra_samples)], dtype=np.int64
    )
    y = np.concatenate([y, y_extra])

    
    templates = 2 * num_classes * np.random.random((num_classes,) + input_shape)
    x = np.zeros((total_samples,) + input_shape, dtype=np.float32)
    for i in range(total_samples):
        x[i] = templates[y[i]] + np.random.normal(
            loc=0, scale=1.0, size=input_shape
        )

    
    indices = np.arange(total_samples)
    np.random.shuffle(indices)
    x, y = x[indices], y[indices]

    
    x_train, y_train, x_test, y_test = [], [], [], []
    for cls in range(num_classes):
        cls_indices = np.where(y == cls)[0]
        np.random.shuffle(cls_indices)
        train_count = int(train_samples / num_classes)

        x_train.extend(x[cls_indices[:train_count]])
        y_train.extend(y[cls_indices[:train_count]])

        x_test.extend(x[cls_indices[train_count:]])
        y_test.extend(y[cls_indices[train_count:]])

    
    x_train, y_train = np.array(x_train), np.array(y_train)
    x_test, y_test = np.array(x_test), np.array(y_test)

    
    train_indices = np.arange(len(x_train))
    test_indices = np.arange(len(x_test))
    np.random.shuffle(train_indices)
    np.random.shuffle(test_indices)

    x_train, y_train = x_train[train_indices], y_train[train_indices]
    x_test, y_test = x_test[test_indices], y_test[test_indices]

    return (x_train, y_train), (x_test, y_test)


def named_product(*args, **kwargs):
    

    def value_to_str(value):
        if hasattr(value, "__name__"):
            return value.__name__.lower()
        return str(value).lower()

    
    all_test_dicts = args + tuple(
        tuple({"testcase_name": value_to_str(v), key: v} for v in values)
        for key, values in kwargs.items()
    )

    
    tests = [{}]
    for test_dicts in all_test_dicts:
        new_tests = []
        for test_dict in test_dicts:
            for test in tests:
                
                testcase_name = test.get("testcase_name", "")
                testcase_name += "_" if testcase_name else ""
                testcase_name += test_dict["testcase_name"]
                new_test = test.copy()
                
                new_test.update(test_dict)
                new_test["testcase_name"] = testcase_name
                new_tests.append(new_test)
        
        tests = new_tests

    return tests