How to use the neuraxle.steps.misc.FitTransformCallbackStep function in neuraxle

def choose_one_step_single_step_chosen_fit():
    a_callback = TapeCallbackFunction()
    b_callback = TapeCallbackFunction()
    c_callback = TapeCallbackFunction()
    d_callback = TapeCallbackFunction()

    return NeuraxleTestCase(
        pipeline=Pipeline([
            ChooseOneStepOf([
                ('a', FitTransformCallbackStep(a_callback, c_callback, transform_function=lambda di: di * 2)),
                ('b', FitTransformCallbackStep(b_callback, d_callback, transform_function=lambda di: di * 2))
            ]),
        ]),
        callbacks=[a_callback, c_callback, b_callback, d_callback],
        expected_callbacks_data=[
            [],
            (DATA_INPUTS, EXPECTED_OUTPUTS),
            [],
            []
        ],
        hyperparams={
            'ChooseOneOrManyStepsOf__choice': 'a'
        },
        expected_processed_outputs=np.array([0, 2, 4, 6, 8, 10, 12, 14, 16, 18])
    )

def choose_one_step_single_step_chosen_fit():
    a_callback = TapeCallbackFunction()
    b_callback = TapeCallbackFunction()
    c_callback = TapeCallbackFunction()
    d_callback = TapeCallbackFunction()

    return NeuraxleTestCase(
        pipeline=Pipeline([
            ChooseOneStepOf([
                ('a', FitTransformCallbackStep(a_callback, c_callback, transform_function=lambda di: di * 2)),
                ('b', FitTransformCallbackStep(b_callback, d_callback, transform_function=lambda di: di * 2))
            ]),
        ]),
        callbacks=[a_callback, c_callback, b_callback, d_callback],
        expected_callbacks_data=[
            [],
            (DATA_INPUTS, EXPECTED_OUTPUTS),
            [],
            []
        ],
        hyperparams={
            'ChooseOneOrManyStepsOf__choice': 'a'
        },
        expected_processed_outputs=np.array([0, 2, 4, 6, 8, 10, 12, 14, 16, 18])
    )

def choose_one_step_single_step_chosen_transform():
    a_callback = TapeCallbackFunction()
    b_callback = TapeCallbackFunction()
    c_callback = TapeCallbackFunction()
    d_callback = TapeCallbackFunction()

    return NeuraxleTestCase(
        pipeline=Pipeline([
            ChooseOneStepOf([
                ('a', FitTransformCallbackStep(a_callback, c_callback, transform_function=lambda di: di * 2)),
                ('b', FitTransformCallbackStep(b_callback, d_callback, transform_function=lambda di: di * 2))
            ]),
        ]),
        callbacks=[a_callback, c_callback, b_callback, d_callback],
        expected_callbacks_data=[
            DATA_INPUTS,
            [],
            [],
            []
        ],
        hyperparams={
            'ChooseOneOrManyStepsOf__choice': 'a'
        },
        expected_processed_outputs=np.array([0, 2, 4, 6, 8, 10, 12, 14, 16, 18])
    )

def create_test_cases():
    data_inputs = np.ones((1, 1))
    expected_outputs = np.ones((1, 1))
    dc = DataContainer(current_ids=range(len(data_inputs)), data_inputs=data_inputs, expected_outputs=expected_outputs)

    tape = TapeCallbackFunction()
    tape_fit = TapeCallbackFunction()
    tape_without_checkpoint_test_arguments = ResumablePipelineTestCase(
        tape,
        data_inputs,
        expected_outputs,
        [
            ("a", FitTransformCallbackStep(tape.callback, tape_fit.callback, ["1"])),
            ("b", FitTransformCallbackStep(tape.callback, tape_fit.callback, ["2"])),
            ("c", FitTransformCallbackStep(tape.callback, tape_fit.callback, ["3"]))
        ],
        ["1", "2", "3"])

    tape2 = TapeCallbackFunction()
    tape2_fit = TapeCallbackFunction()
    tape_checkpoint_not_saved_test_arguments = ResumablePipelineTestCase(
        tape2,
        data_inputs,
        expected_outputs,
        [
            ("a", FitTransformCallbackStep(tape2.callback, tape2_fit.callback, ["1"])),
            ("b", SomeCheckpointStep(data_container=None)
             ),
            ("c", FitTransformCallbackStep(tape2.callback, tape2_fit.callback, ["2"])),

def create_checkpoint_test_case(tmpdir):
    tape_transform_1 = TapeCallbackFunction()
    tape_fit_1 = TapeCallbackFunction()
    tape_transform_2 = TapeCallbackFunction()
    tape_fit_2 = TapeCallbackFunction()
    pipeline = ResumablePipeline([
        ('step1', FitTransformCallbackStep(tape_transform_1, tape_fit_1)),
        ('checkpoint', DefaultCheckpoint()),
        ('step2', FitTransformCallbackStep(tape_transform_2, tape_fit_2))
    ], cache_folder=tmpdir)

    return CheckpointTest(
        tape_transform_1, tape_fit_1, tape_transform_2, tape_fit_2, pipeline
    )

            EpochRepeater(FitTransformCallbackStep(callback_transform, callback_fit), epochs=EPOCHS)
        ]),
        callbacks=[callback_fit, callback_transform],
        expected_callbacks_data=[
            [(DATA_INPUTS, EXPECTED_OUTPUTS), (DATA_INPUTS, EXPECTED_OUTPUTS)],
            []
        ],
        data_inputs=DATA_INPUTS,
        expected_outputs=EXPECTED_OUTPUTS,
        execution_mode=ExecutionMode.FIT
    )
])
def test_epoch_repeater(test_case):
    processed_outputs = test_case.execute()

    test_case.assert_expected_processed_outputs(processed_outputs)
    test_case.assert_callback_data_is_as_expected()

def create_full_execution_checkpoint_test_case(tmpdir):
    tape_transform_1 = TapeCallbackFunction()
    tape_fit_1 = TapeCallbackFunction()
    tape_transform_2 = TapeCallbackFunction()
    tape_fit_2 = TapeCallbackFunction()
    pipeline = ResumablePipeline([
        ('step1', FitTransformCallbackStep(tape_transform_1, tape_fit_1)),
        ('checkpoint', FullExecutionCheckpoint()),
        ('step2', FitTransformCallbackStep(tape_transform_2, tape_fit_2))
    ], cache_folder=tmpdir)

    return CheckpointTest(
        tape_transform_1, tape_fit_1, tape_transform_2, tape_fit_2, pipeline
    )

["3", "4"])

    tape7 = TapeCallbackFunction()
    tape7_fit = TapeCallbackFunction()
    tape_checkpoint_saved_inside_subpipeline_first_step = ResumablePipelineTestCase(
        tape7,
        data_inputs,
        expected_outputs,
        [
            ("a", FitTransformCallbackStep(tape7.callback, tape7_fit.callback, ["1"])),
            ResumablePipeline([
                ("d", SomeCheckpointStep(data_container=dc)
                 ),
                ("b", FitTransformCallbackStep(tape7.callback, tape7_fit.callback, ["2"])),
            ]),
            ("e", FitTransformCallbackStep(tape7.callback, tape7_fit.callback, ["3"])),
            ("f", FitTransformCallbackStep(tape7.callback, tape7_fit.callback, ["4"])),
        ],
        ["2", "3", "4"])

    tape8 = TapeCallbackFunction()
    tape8_fit = TapeCallbackFunction()
    tape_checkpoint_saved_inside_subpipeline_step_in_the_middle = ResumablePipelineTestCase(
        tape8,
        data_inputs,
        expected_outputs,
        [
            ("a", FitTransformCallbackStep(tape8.callback, tape8_fit.callback, ["1"])),
            ResumablePipeline([
                ("b", FitTransformCallbackStep(tape8.callback, tape8_fit.callback, ["2"])),
                ("d", SomeCheckpointStep(data_container=dc)
                 ),

def create_test_case_fit_single_step_choosen():
    a_callback = TapeCallbackFunction()
    b_callback = TapeCallbackFunction()
    c_callback = TapeCallbackFunction()
    d_callback = TapeCallbackFunction()

    return NeuraxleTestCase(
        pipeline=Pipeline([
            ChooseOneOrManyStepsOf([
                ('a', FitTransformCallbackStep(a_callback, c_callback, transform_function=lambda di: di * 2)),
                ('b', FitTransformCallbackStep(b_callback, d_callback, transform_function=lambda di: di * 2))
            ]),
        ]),
        callbacks=[a_callback, c_callback, b_callback, d_callback],
        expected_callbacks_data=[
            [],
            (DATA_INPUTS, EXPECTED_OUTPUTS),
            [],
            []
        ],
        hyperparams={
            'ChooseOneOrManyStepsOf__a__enabled': True,
            'ChooseOneOrManyStepsOf__b__enabled': False
        },
        hyperparams_space={
            'ChooseOneOrManyStepsOf__a__enabled': Boolean(),
            'ChooseOneOrManyStepsOf__b__enabled': Boolean()