AutoGrad_CPP/FuncWrapper.h at main · gavinchen430/AutoGrad_CPP · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
#include <functional>
#include "Box.h"

template<typename F>
struct FunctionTraits;

template<typename ClassType, typename R, typename... Args>
struct FunctionTraits<R(ClassType::*)(Args...)>
{
    FunctionTraits() {}
    using RetType = R;
    using ArgTypes = std::tuple<Args...>;
    ClassType mCT;
    FunctionTraits(ClassType& ct) : mCT(ct) {
    }

    R operator()(Args... args) {
        auto r = mCT(args.incTraceID()...);
        return backward_pass<R>(r.decTraceID());
    }
};

template<typename F>
struct FunctionTraits : FunctionTraits<decltype(&F::operator())> {
    FunctionTraits(F& f) : FunctionTraits<decltype(&F::operator())>(f){
    }
};


template<typename R, typename... Args>
struct FunctionTraits<R(*)(Args...)>
{
    FunctionTraits() {
    }
    using Pointer = R(*)(Args...);
    Pointer mP;
    FunctionTraits(Pointer p) : mP(p) {
    }

    R operator()(Args... args) {
        auto r = mP(args.incTraceID()...);
        return backward_pass<R>(r.decTraceID());
    }
};

template<typename T>
FunctionTraits<T> grad(T t) {
    return FunctionTraits<T>(t);
}