aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/framework/function_testlib.h
blob: ed0446ea857d2beff123d84ed037cb3e6bc33b91 (plain)
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
50
51
52
53
#ifndef TENSORFLOW_FRAMEWORK_FUNCTION_TESTLIB_H_
#define TENSORFLOW_FRAMEWORK_FUNCTION_TESTLIB_H_

#include <string>

#include "tensorflow/core/framework/attr_value_util.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
#include "tensorflow/core/platform/port.h"

namespace tensorflow {
namespace test {
namespace function {

// Helper to construct a NodeDef.
NodeDef NDef(
    const string& name, const string& op, gtl::ArraySlice<string> inputs,
    gtl::ArraySlice<std::pair<string, FunctionDefHelper::AttrValueWrapper>>
        attrs = {},
    const string& device = "");

// Helper to construct a GraphDef proto.
GraphDef GDef(gtl::ArraySlice<NodeDef> nodes,
              gtl::ArraySlice<FunctionDef> funcs = {});

// For testing convenience, we provide a few simple functions that can
// be easily executed and tested.

// x:T -> x * 2.
FunctionDef XTimesTwo();

// x:T -> (x * 2) * 2.
FunctionDef XTimesFour();

// x:T -> ((x * 2) * 2) * 2.
FunctionDef XTimes16();

// w:T, x:T, b:T -> MatMul(w, x) + b
FunctionDef WXPlusB();

// x:T -> x:T, T is a type which we automatically converts to a bool.
FunctionDef NonZero();

// x:T, y:T -> y:T, x:T
FunctionDef Swap();

}  // end namespace function
}  // end namespace test
}  // end namespace tensorflow

#endif  // TENSORFLOW_FRAMEWORK_FUNCTION_TESTLIB_H_