test_vm.py 文件源码

def test_partial_function_with_output_keys():

    def check_partial_function_output_keys(linker_name):
        x = tensor.scalar('input')
        y = 3 * x
        f = theano.function([x], {'a': y * 5, 'b': y - 7}, mode=Mode(
            optimizer=None, linker=linker_name))

        assert f(5, output_subset=['a'])['a'] == f(5)['a']

    check_partial_function_output_keys(vm.VM_Linker(allow_partial_eval=True, use_cloop=False))
    check_partial_function_output_keys('cvm')