def test_tcn_forecaster_xshard_input(self):
    (train_data, val_data, test_data) = create_data()
    print('original', train_data[0].dtype)
    init_orca_context(cores=4, memory='2g')
    from zoo.orca.data import XShards

    def transform_to_dict(data):
        return {'x': data[0], 'y': data[1]}

    def transform_to_dict_x(data):
        return {'x': data[0]}
    train_data = XShards.partition(train_data).transform_shard(transform_to_dict)
    val_data = XShards.partition(val_data).transform_shard(transform_to_dict)
    test_data = XShards.partition(test_data).transform_shard(transform_to_dict_x)
    for distributed in [True, False]:
        forecaster = LSTMForecaster(past_seq_len=24, input_feature_num=2, output_feature_num=2, loss='mae', lr=0.01, distributed=distributed)
        forecaster.fit(train_data, epochs=2)
        distributed_pred = forecaster.predict(test_data)
        distributed_eval = forecaster.evaluate(val_data)
    stop_orca_context()
def test_tcn_forecaster_xshard_input(self):
    (train_data, val_data, test_data) = create_data()
    print('original', train_data[0].dtype)
    init_orca_context(cores=4, memory='2g')
    from zoo.orca.data import XShards

    def transform_to_dict(data):
        return {'x': data[0], 'y': data[1]}

    def transform_to_dict_x(data):
        return {'x': data[0]}
    train_data = XShards.partition(train_data).transform_shard(transform_to_dict)
    val_data = XShards.partition(val_data).transform_shard(transform_to_dict)
    test_data = XShards.partition(test_data).transform_shard(transform_to_dict_x)
    for distributed in [True, False]:
        forecaster = LSTMForecaster(past_seq_len=24, input_feature_num=2, output_feature_num=2, loss='mae', lr=0.01, distributed=distributed)
        forecaster.fit(train_data, epochs=2)
        distributed_pred = forecaster.predict(test_data)
        distributed_eval = forecaster.evaluate(val_data)
    stop_orca_context()