def get_correctness_test_inputs(use_numpy, with_distribution, x_train, y_train, x_predict):
    'Generates the inputs for correctness check when enable Keras with DS.'
    global_batch_size = 64
    batch_size = global_batch_size
    use_per_core_batch_size = (with_distribution and (with_distribution.__class__.__name__ != 'TPUStrategy'))
    if use_per_core_batch_size:
        batch_size //= with_distribution.num_replicas_in_sync
    if use_numpy:
        training_inputs = {'batch_size': batch_size, 'x': x_train, 'y': y_train, 'epochs': 1, 'shuffle': False}
        eval_inputs = {'batch_size': batch_size, 'x': x_train, 'y': y_train}
        if with_distribution:
            if use_per_core_batch_size:
                predict_batch_size = (len(x_predict) // with_distribution.num_replicas_in_sync)
            else:
                predict_batch_size = len(x_predict)
        else:
            predict_batch_size = None
        predict_inputs = {'batch_size': predict_batch_size, 'x': np.array(x_predict, dtype=np.float32)}
    else:
        train_dataset = dataset_ops.Dataset.from_tensor_slices((x_train, y_train))
        x = batch_wrapper(train_dataset, batch_size, with_distribution)
        training_inputs = {'batch_size': None, 'x': x, 'y': None, 'epochs': 1, 'shuffle': False, 'steps_per_epoch': (len(x_train) // global_batch_size)}
        eval_inputs = {'batch_size': None, 'x': x, 'y': None, 'steps': 20}
        predict_batch_size = len(x_predict)
        if use_per_core_batch_size:
            predict_batch_size //= with_distribution.num_replicas_in_sync
        predict_dataset = dataset_ops.Dataset.from_tensor_slices(x_predict)
        predict_dataset = batch_wrapper(predict_dataset, predict_batch_size, with_distribution)
        predict_inputs = {'batch_size': None, 'steps': 1, 'x': predict_dataset}
    return (training_inputs, eval_inputs, predict_inputs)
Generates the inputs for correctness check when enable Keras with DS.
def get_correctness_test_inputs(use_numpy, with_distribution, x_train, y_train, x_predict):
    
    global_batch_size = 64
    batch_size = global_batch_size
    use_per_core_batch_size = (with_distribution and (with_distribution.__class__.__name__ != 'TPUStrategy'))
    if use_per_core_batch_size:
        batch_size //= with_distribution.num_replicas_in_sync
    if use_numpy:
        training_inputs = {'batch_size': batch_size, 'x': x_train, 'y': y_train, 'epochs': 1, 'shuffle': False}
        eval_inputs = {'batch_size': batch_size, 'x': x_train, 'y': y_train}
        if with_distribution:
            if use_per_core_batch_size:
                predict_batch_size = (len(x_predict) // with_distribution.num_replicas_in_sync)
            else:
                predict_batch_size = len(x_predict)
        else:
            predict_batch_size = None
        predict_inputs = {'batch_size': predict_batch_size, 'x': np.array(x_predict, dtype=np.float32)}
    else:
        train_dataset = dataset_ops.Dataset.from_tensor_slices((x_train, y_train))
        x = batch_wrapper(train_dataset, batch_size, with_distribution)
        training_inputs = {'batch_size': None, 'x': x, 'y': None, 'epochs': 1, 'shuffle': False, 'steps_per_epoch': (len(x_train) // global_batch_size)}
        eval_inputs = {'batch_size': None, 'x': x, 'y': None, 'steps': 20}
        predict_batch_size = len(x_predict)
        if use_per_core_batch_size:
            predict_batch_size //= with_distribution.num_replicas_in_sync
        predict_dataset = dataset_ops.Dataset.from_tensor_slices(x_predict)
        predict_dataset = batch_wrapper(predict_dataset, predict_batch_size, with_distribution)
        predict_inputs = {'batch_size': None, 'steps': 1, 'x': predict_dataset}
    return (training_inputs, eval_inputs, predict_inputs)