Function TestDataset

tensorflow/go/op/op_test.go:96–133

(t *testing.T)

}

func TestDataset(t *testing.T) {
	var (
		s = NewScope()

		// The use of a non-scalar here is inspired by
		// https://github.com/tensorflow/tensorflow/issues/14891
		c       = Const(s, []int32{21718, 31415})
		types   = []tf.DataType{c.DataType()}
		shapes  = []tf.Shape{c.Shape()}
		dataset = TensorDataset(s, []tf.Output{c}, shapes)

		iterator = Iterator(s, "", "", types, shapes)
		next     = IteratorGetNext(s, iterator, types, shapes)
		init     = MakeIterator(s, dataset, iterator)
	)
	graph, err := s.Finalize()
	if err != nil {
		t.Fatal(err)
	}
	sess, err := tf.NewSession(graph, nil)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := sess.Run(nil, nil, []*tf.Operation{init}); err != nil {
		t.Fatal(err)
	}
	results, err := sess.Run(nil, next, nil)
	if err != nil {
		t.Fatal(err)
	}
	got := results[0].Value().([]int32)
	if len(got) != 2 || got[0] != 21718 || got[1] != 31415 {
		t.Errorf("Got %v, want {21718, 31415}", got)
	}
	if _, err := sess.Run(nil, next, nil); err == nil {
		t.Errorf("Expected sess.Run() to fail since the iterator should have reached the end of the dataset")
	}
}