Class OpTestBuilder

tensorflow/compiler/tests/randomized_tests.cc:96–146

An OpTestBuilder is a graph builder class that takes as input an operator to test, its inputs and attributes, and builds a graph that executes the operator.

// test, its inputs and attributes, and builds a graph that executes the
// operator.
class OpTestBuilder {
 public:
  explicit OpTestBuilder(const string& op_name);

  // Adds an input 'tensor' as a Placeholder node.
  OpTestBuilder& Input(const Tensor& tensor);

  // Adds a random input tensor with 'type' as a Placeholder node.
  // If 'dims' is not provided, RandomDims() is used.
  OpTestBuilder& RandomInput(DataType type);
  OpTestBuilder& RandomInput(DataType type, std::vector<int64> dims);

  // As RandomInput but the values are unique.
  OpTestBuilder& RandomUniqueInput(DataType type, std::vector<int64> dims);

  // Sets an attribute.
  template <class T>
  OpTestBuilder& Attr(absl::string_view attr_name, T&& value);

  // Overload needed to allow {...} expressions for value.
  template <class T>
  OpTestBuilder& Attr(absl::string_view attr_name,
                      std::initializer_list<T> value);

  // Adds nodes that executes the operator under test on 'device' to 'graphdef'.
  // If 'use_jit' is true, marks the operator under test to be compiled by XLA.
  // The graph will consist of one Placeholder node per input, the operator
  // itself, and one Identity node per output. If 'test_node_def' is not null,
  // sets it to the NodeDef of the operator under test. Fills 'inputs' and
  // 'outputs' with the names of the input placeholder nodes and the output
  // identity nodes, respectively.
  Status BuildGraph(const string& name_prefix, const string& device,
                    bool use_jit, GraphDef* graphdef, NodeDef** test_node_def,
                    std::vector<string>* inputs,
                    std::vector<string>* outputs) const;

  struct InputDescription {
    Tensor tensor;

    DataType type = DT_INVALID;
    bool has_dims = false;
    bool needs_unique_values = false;
    std::vector<int64> dims;
  };

  const std::vector<InputDescription>& inputs() const { return inputs_; }

 private:
  NodeDef node_def_;
  std::vector<InputDescription> inputs_;
};

OpTestBuilder::OpTestBuilder(const string& op_name) {
  node_def_.set_op(op_name);