/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// Randomized tests for XLA implementations of Tensorflow operations.
//
// For each operator, the tests in this file choose a set of random inputs and
// attributes. The test then compares the outputs of the operator when executed
// via Tensorflow using the CPU device and when executed via XLA.
//
// By default, each test chooses a random seed nondeterministically (using
// std::random_device). However, a particular choice of random seed can be
// forced using the flag --tf_xla_random_seed; each test logs the
// flag value necessary to reproduce its outputs.
//
// Example usage:
// Run tests, comparing the Tensorflow CPU operators with their XLA-compiled
// counterparts:
// randomized_tests \
//   --tf_xla_test_use_jit=true --tf_xla_test_device=CPU:0 \
//   --tf_xla_test_repetitions=20

// TODO(phawkins): add tests for:
// * ArgMax
// * DepthwiseConv2DNative
// * Gather
// * InvertPermutation
// * MaxPoolGrad (requires implementation of forward operator)
// * Select
// * Unpack
//
// TODO(phawkins): improve tests for:
// * StridedSliceGrad (need to use shape function to compute sensible inputs)

#include <random>
#include <unordered_map>

#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/tf2xla/type_util.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/graph_constructor.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/public/session_options.h"
#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/core/util/device_name_utils.h"
#include "tensorflow/core/util/tensor_format.h"

namespace tensorflow {
namespace {

// Command line flags: see main() below.
int64 tf_xla_random_seed = 0;
int32 tf_xla_test_repetitions = 20;
string* tf_xla_test_device_ptr;  // initial value set in main()
bool tf_xla_test_use_jit = true;

string LocalDeviceToFullDeviceName(const string& device) {
  return strings::StrCat("/job:localhost/replica:0/task:0/device:", device);
}

constexpr std::array<DataType, 3> kAllXlaTypes = {
    {DT_INT32, DT_FLOAT, DT_BOOL}};

// 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.
class OpTestBuilder {
 public:
  explicit OpTestBuilder(const string& op_name);

  // Adds an input 'tensor'.
  OpTestBuilder& Input(const Tensor& tensor);

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

  // Overload needed to allow {...} expressions for value.
  template <class T>
  OpTestBuilder& Attr(StringPiece 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;

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

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

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

OpTestBuilder& OpTestBuilder::Input(const Tensor& tensor) {
  VLOG(1) << "Adding input: " << tensor.DebugString();
  inputs_.push_back(tensor);
  return *this;
}

template <class T>
OpTestBuilder& OpTestBuilder::Attr(StringPiece attr_name, T&& value) {
  AddNodeAttr(attr_name, std::forward<T>(value), &node_def_);
  return *this;
}

template <class T>
OpTestBuilder& OpTestBuilder::Attr(StringPiece attr_name,
                                   std::initializer_list<T> value) {
  Attr<std::initializer_list<T>>(attr_name, std::move(value));
  return *this;
}

Status OpTestBuilder::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 {
  OpRegistryInterface* op_registry = OpRegistry::Global();

  const OpDef* op_def;
  TF_RETURN_IF_ERROR(op_registry->LookUpOpDef(node_def_.op(), &op_def));

  NodeDef* test_def = graphdef->add_node();
  *test_def = node_def_;
  test_def->set_name(strings::StrCat(name_prefix, "_op_under_test"));
  test_def->set_device(device);
  AddDefaultsToNodeDef(*op_def, test_def);
  if (use_jit) {
    AddNodeAttr(kXlaCompileAttr, true, test_def);
  }
  VLOG(1) << "Op under test: " << test_def->DebugString();

  DataTypeVector input_types, output_types;
  TF_RETURN_IF_ERROR(
      InOutTypesForNode(*test_def, *op_def, &input_types, &output_types));

  // Build feed and fetch nodes.
  for (int i = 0; i < input_types.size(); ++i) {
    NodeDef* def = graphdef->add_node();
    string name = strings::StrCat(name_prefix, "_input_", i);
    TF_RETURN_IF_ERROR(NodeDefBuilder(name, "Placeholder")
                           .Device(device)
                           .Attr("dtype", input_types[i])
                           .Finalize(def));
    inputs->push_back(name);
    test_def->add_input(name);
  }

  for (int i = 0; i < output_types.size(); ++i) {
    NodeDef* def = graphdef->add_node();
    string name = strings::StrCat(name_prefix, "_output_", i);
    TF_RETURN_IF_ERROR(NodeDefBuilder(name, "Identity")
                           .Device(device)
                           .Attr("T", output_types[i])
                           .Input(test_def->name(), i, output_types[i])
                           .Finalize(def));
    outputs->push_back(name);
  }

  if (test_node_def) {
    *test_node_def = test_def;
  }

  return Status::OK();
}

// Test fixture. The fixture manages the random number generator and its seed,
// and has a number of convenience methods for building random Tensors, shapes,
// etc.
class OpTest : public ::testing::Test {
 public:
  OpTest();

  // Runs 'fn' up to --tf_xla_test_repetitions times, or until a failure occurs;
  // whichever happens first.
  void Repeatedly(const std::function<void(void)>& fn);

  // Select a random element from 'candidates'.
  template <typename T>
  T Choose(gtl::ArraySlice<T> candidates);

  static constexpr int kDefaultMaxRank = 5;
  static constexpr int64 kDefaultMaxDimensionSize = 20LL;

  // Returns a random dimension size, in the range [min, max).
  int64 RandomDim(int64 min = 0, int64 max = kDefaultMaxDimensionSize);

  // Returns a random shape. The tensor has rank in the range [min_rank,
  // max_rank). Each dimension has size [min_size, max_size).
  std::vector<int64> RandomDims(int min_rank = 0,
                                int max_rank = kDefaultMaxRank,
                                int64 min_size = 0,
                                int64 max_size = kDefaultMaxDimensionSize);

  // Given a shape 'dims', build a pair of dimensions such that one broadcasts
  // to the other.
  std::pair<std::vector<int64>, std::vector<int64>> BroadcastableDims(
      std::vector<int64> dims);

  // Builds a random pair of broadcastable dims.
  // TODO(phawkins): currently the maximum rank is 3, because broadcasting > 3
  // dimensions is unimplemented by the Tensorflow Eigen code (b/29268487)
  std::pair<std::vector<int64>, std::vector<int64>> BroadcastableDims();

  // Returns a tensor filled with random but "reasonable" values from the middle
  // of the type's range. If the shape is omitted, a random shape is used.
  // TODO(phawkins): generalize this code to a caller-supplied distribution.
  Tensor RandomTensor(DataType dtype, gtl::ArraySlice<int64> shape);
  Tensor RandomTensor(DataType dtype);

  // Like RandomTensor, but uses values >= 0.
  Tensor RandomNonNegativeTensor(DataType dtype, gtl::ArraySlice<int64> shape);
  Tensor RandomNonNegativeTensor(DataType dtype);

  // Returns a random subset of the integers in the range [0, rank), suitable
  // for use as reduction indices.
  Tensor RandomReductionIndices(int rank);

  struct WindowedSpatialDims {
    Padding padding;
    std::vector<int64> kernel_dims;
    std::vector<int64> stride_dims;
    std::vector<int64> input_dims;
    std::vector<int64> output_dims;
  };
  // Choose spatial dimensions for a windowed op such as pooling or convolution.
  WindowedSpatialDims ChooseWindowedSpatialDims(int num_spatial_dims);

  // Builds dimensions for a windowed op such as pooling or convolution,
  // including a batch and feature dimension.
  std::vector<int64> ImageDims(TensorFormat format, int batch, int feature,
                               const std::vector<int64>& spatial_dims);

  // Converts an int64 vector to an int32 vector.
  std::vector<int32> AsInt32s(const std::vector<int64>& int64s);

  std::mt19937& generator() { return *generator_; }

  // Run the test case described by 'builder' with and without XLA and check
  // that the outputs are close. Tensors x and y are close if they have the same
  // type, same shape, and have close values. For floating-point tensors, the
  // element-wise difference between x and y must no more than
  // atol + rtol * abs(x); or both elements may be NaN or infinity. For
  // non-floating-point tensors the element values must match exactly.
  void ExpectTfAndXlaOutputsAreClose(const OpTestBuilder& builder,
                                     double atol = 1e-2, double rtol = 1e-2);

 protected:
  // Per-test state:
  std::unique_ptr<std::mt19937> generator_;

  std::unique_ptr<Session> session_;

  // Number of test cases built in 'session_'. Used to uniquify node names.
  int num_tests_ = 0;
};

OpTest::OpTest() {
  // Creates a random-number generator for the test case. Use the value of
  // --tf_xla_random_seed as the seed, if provided.
  int64 s = tf_xla_random_seed;
  unsigned int seed;
  if (s <= 0) {
    std::random_device random_device;
    seed = random_device();
  } else {
    seed = static_cast<unsigned int>(s);
  }
  LOG(INFO) << "Random seed for test case: " << seed
            << ". To reproduce the "
               "results of this test, pass flag --tf_xla_random_seed="
            << seed;
  generator_.reset(new std::mt19937(seed));

  // Create a session with an empty graph.
  SessionOptions session_options;
  session_.reset(NewSession(session_options));
  GraphDef def;
  TF_CHECK_OK(session_->Create(def));
}

void OpTest::Repeatedly(const std::function<void(void)>& fn) {
  int const max_repetitions = tf_xla_test_repetitions;
  for (int i = 0; !HasFailure() && i < max_repetitions; ++i) {
    fn();
  }
}

template <typename T>
T OpTest::Choose(gtl::ArraySlice<T> candidates) {
  std::uniform_int_distribution<size_t> d(0, candidates.size() - 1);
  return candidates[d(generator())];
}

int64 OpTest::RandomDim(int64 min, int64 max) {
  std::uniform_int_distribution<int64> size_distribution(min, max - 1);
  return size_distribution(generator());
}

std::vector<int64> OpTest::RandomDims(int min_rank, int max_rank,
                                      int64 min_size, int64 max_size) {
  CHECK_LE(0, min_rank);
  CHECK_LE(min_rank, max_rank);
  std::uniform_int_distribution<int> rank_distribution(min_rank, max_rank);
  int rank = rank_distribution(generator());
  std::vector<int64> dims(rank);
  std::generate(dims.begin(), dims.end(), [this, min_size, max_size]() {
    return RandomDim(min_size, max_size);
  });
  return dims;
}

Tensor OpTest::RandomTensor(DataType dtype, gtl::ArraySlice<int64> shape) {
  Tensor tensor(dtype, TensorShape(shape));
  switch (dtype) {
    case DT_FLOAT: {
      std::uniform_real_distribution<float> distribution(-1.0f, 1.0f);
      test::FillFn<float>(&tensor, [this, &distribution](int i) -> float {
        return distribution(generator());
      });
      break;
    }
    case DT_DOUBLE: {
      std::uniform_real_distribution<double> distribution(-1.0, 1.0);
      test::FillFn<double>(&tensor, [this, &distribution](int i) -> double {
        return distribution(generator());
      });
      break;
    }
    case DT_INT32: {
      std::uniform_int_distribution<int32> distribution(-(1 << 20), 1 << 20);
      test::FillFn<int32>(&tensor, [this, &distribution](int i) -> int32 {
        return distribution(generator());
      });
      break;
    }
    case DT_INT64: {
      std::uniform_int_distribution<int64> distribution(-(1LL << 40),
                                                        1LL << 40);
      test::FillFn<int64>(&tensor, [this, &distribution](int i) -> int64 {
        return distribution(generator());
      });
      break;
    }
    case DT_BOOL: {
      std::bernoulli_distribution distribution;
      test::FillFn<bool>(&tensor, [this, &distribution](int i) -> bool {
        return distribution(generator());
      });
      break;
    }
    default:
      LOG(FATAL) << "Unimplemented type " << dtype << " in RandomTensor";
  }
  return tensor;
}

Tensor OpTest::RandomTensor(DataType dtype) {
  return RandomTensor(dtype, RandomDims());
}

Tensor OpTest::RandomNonNegativeTensor(DataType dtype,
                                       gtl::ArraySlice<int64> shape) {
  Tensor tensor(dtype, TensorShape(shape));
  switch (dtype) {
    case DT_FLOAT: {
      std::uniform_real_distribution<float> distribution(0.0f, 1.0f);
      test::FillFn<float>(&tensor, [this, &distribution](int i) -> float {
        return distribution(generator());
      });
      break;
    }
    case DT_DOUBLE: {
      std::uniform_real_distribution<double> distribution(0.0, 1.0);
      test::FillFn<double>(&tensor, [this, &distribution](int i) -> double {
        return distribution(generator());
      });
      break;
    }
    case DT_INT32: {
      std::uniform_int_distribution<int32> distribution(0, 1 << 20);
      test::FillFn<int32>(&tensor, [this, &distribution](int i) -> int32 {
        return distribution(generator());
      });
      break;
    }
    case DT_INT64: {
      std::uniform_int_distribution<int64> distribution(0, 1LL << 40);
      test::FillFn<int64>(&tensor, [this, &distribution](int i) -> int64 {
        return distribution(generator());
      });
      break;
    }
    default:
      LOG(FATAL) << "Unimplemented type " << dtype
                 << " in RandomNonNegativeTensor";
  }
  return tensor;
}

Tensor OpTest::RandomNonNegativeTensor(DataType dtype) {
  return RandomNonNegativeTensor(dtype, RandomDims());
}

std::pair<std::vector<int64>, std::vector<int64>> OpTest::BroadcastableDims(
    std::vector<int64> dims) {
  if (dims.empty()) return {dims, dims};

  // Remove some dimensions from the front of 'dims'.
  size_t skip =
      std::uniform_int_distribution<size_t>(0, dims.size() - 1)(generator());

  std::vector<int64> bdims(dims.begin() + skip, dims.end());

  // Randomly replace some of the remaining dimensions of 'dims' with 1.
  std::bernoulli_distribution random_bool;

  for (int64& dim : bdims) {
    if (random_bool(generator())) {
      dim = 1LL;
    }
  }

  // Possibly swap the roles of 'dims' and 'bdims'.
  if (random_bool(generator())) {
    dims.swap(bdims);
  }
  return {dims, bdims};
}

std::pair<std::vector<int64>, std::vector<int64>> OpTest::BroadcastableDims() {
  return BroadcastableDims(RandomDims(0, 3));
}

Tensor OpTest::RandomReductionIndices(int rank) {
  std::bernoulli_distribution random_bool;
  std::vector<int32> indices;
  for (int i = 0; i < rank; ++i) {
    if (random_bool(generator())) {
      indices.push_back(i);
    }
  }
  return test::AsTensor<int32>(indices);
}

OpTest::WindowedSpatialDims OpTest::ChooseWindowedSpatialDims(
    int num_spatial_dims) {
  WindowedSpatialDims d;
  d.padding = Choose<Padding>({SAME, VALID});
  std::uniform_int_distribution<int> random_int(1, 5);
  d.kernel_dims.resize(num_spatial_dims);
  d.input_dims.resize(num_spatial_dims);
  d.output_dims.resize(num_spatial_dims);
  d.stride_dims.resize(num_spatial_dims);
  for (int i = 0; i < num_spatial_dims; ++i) {
    Status s;
    // Repeatedly try different filter/stride sizes until we find a valid
    // combination.
    do {
      // CPU implementations require stride <= kernel size.
      d.kernel_dims[i] = random_int(generator()),
      d.input_dims[i] = RandomDim(d.kernel_dims[i]);
      d.stride_dims[i] =
          std::uniform_int_distribution<int>(1, d.kernel_dims[i])(generator());
      int64 pad_dummy;
      s = GetWindowedOutputSize(d.input_dims[i], d.kernel_dims[i],
                                d.stride_dims[i], d.padding, &d.output_dims[i],
                                &pad_dummy);
    } while (!s.ok());
  }
  return d;
}

std::vector<int64> OpTest::ImageDims(TensorFormat format, int batch,
                                     int feature,
                                     const std::vector<int64>& spatial_dims) {
  std::vector<int64> dims;
  switch (format) {
    case FORMAT_NHWC:
      dims.push_back(batch);
      for (int dim : spatial_dims) {
        dims.push_back(dim);
      }
      dims.push_back(feature);
      break;
    case FORMAT_NCHW:
      dims.push_back(batch);
      dims.push_back(feature);
      for (int dim : spatial_dims) {
        dims.push_back(dim);
      }
      break;
  }
  return dims;
}

std::vector<int32> OpTest::AsInt32s(const std::vector<int64>& int64s) {
  return std::vector<int32>(int64s.begin(), int64s.end());
}

// Functions for comparing tensors.

template <typename T>
bool IsClose(const T& x, const T& y, double atol, double rtol) {
  if (std::isnan(x) && std::isnan(y)) return true;
  if (x == y) return true;  // Allow inf == inf.
  return fabs(x - y) < atol + rtol * fabs(x);
}

template <typename T>
Status TensorsAreCloseImpl(const Tensor& x, const Tensor& y, double atol,
                           double rtol) {
  auto Tx = x.flat<T>();
  auto Ty = y.flat<T>();
  for (int i = 0; i < Tx.size(); ++i) {
    if (!IsClose(Tx(i), Ty(i), atol, rtol)) {
      return errors::InvalidArgument(strings::StrCat(
          i, "-th tensor element isn't close: ", Tx(i), " vs. ", Ty(i),
          ". x = ", x.DebugString(), "y = ", y.DebugString(), "atol = ", atol,
          " rtol = ", rtol, " tol = ", atol + rtol * std::fabs(Tx(i))));
    }
  }
  return Status::OK();
}

template <typename T>
Status TensorsAreEqualImpl(const Tensor& x, const Tensor& y) {
  auto Tx = x.flat<T>();
  auto Ty = y.flat<T>();
  for (int i = 0; i < Tx.size(); ++i) {
    if (Tx(i) != Ty(i)) {
      return errors::InvalidArgument(strings::StrCat(
          i, "-th tensor element isn't equal: ", Tx(i), " vs. ", Ty(i),
          ". x = ", x.DebugString(), "y = ", y.DebugString()));
    }
  }
  return Status::OK();
}

// Tests if "x" and "y" are tensors of the same type, same shape, and with
// close values. For floating-point tensors, the element-wise difference between
// x and y must no more than atol + rtol * abs(x). For non-floating-point
// tensors the values must match exactly.
Status TensorsAreClose(const Tensor& a, const Tensor& b, double atol,
                       double rtol) {
  if (a.dtype() != b.dtype()) {
    return errors::InvalidArgument(strings::StrCat(
        "Tensors have different types: ", DataTypeString(a.dtype()), " and ",
        DataTypeString(b.dtype())));
  }
  if (!a.IsSameSize(b)) {
    return errors::InvalidArgument(strings::StrCat(
        "Tensors have different shapes: ", a.shape().DebugString(), " and ",
        b.shape().DebugString()));
  }

  switch (a.dtype()) {
    case DT_FLOAT:
      return TensorsAreCloseImpl<float>(a, b, atol, rtol);
    case DT_DOUBLE:
      return TensorsAreCloseImpl<double>(a, b, atol, rtol);
    case DT_INT32:
      return TensorsAreEqualImpl<int32>(a, b);
    case DT_INT64:
      return TensorsAreEqualImpl<int64>(a, b);
    case DT_BOOL:
      return TensorsAreEqualImpl<bool>(a, b);
    default:
      LOG(FATAL) << "Unexpected type : " << DataTypeString(a.dtype());
  }
}

void OpTest::ExpectTfAndXlaOutputsAreClose(const OpTestBuilder& builder,
                                           double atol, double rtol) {
  string cpu_device =
      LocalDeviceToFullDeviceName(strings::StrCat(DEVICE_CPU, ":0"));
  string test_device = LocalDeviceToFullDeviceName(*tf_xla_test_device_ptr);

  DeviceNameUtils::ParsedName parsed_name;
  ASSERT_TRUE(
      DeviceNameUtils::ParseLocalName(*tf_xla_test_device_ptr, &parsed_name));
  DeviceType test_device_type(parsed_name.type);
  ++num_tests_;

  GraphDef graph;
  std::vector<string> expected_inputs, test_inputs;
  std::vector<string> expected_fetches, test_fetches;
  TF_ASSERT_OK(builder.BuildGraph(
      strings::StrCat("test", num_tests_, "_expected"), cpu_device,
      /* use_jit= */ false, &graph, /* test_node_def= */ nullptr,
      &expected_inputs, &expected_fetches));

  NodeDef* node_def;
  TF_ASSERT_OK(builder.BuildGraph(strings::StrCat("test", num_tests_, "_test"),
                                  test_device, tf_xla_test_use_jit, &graph,
                                  &node_def, &test_inputs, &test_fetches));

  // Check that there's a kernel corresponding to 'node_def' on the device under
  // test.
  Status status = FindKernelDef(test_device_type, *node_def, nullptr, nullptr);
  if (!status.ok()) {
    VLOG(1) << "Skipping test because there is no corresponding registered "
            << "kernel on the test device: " << status;
    return;
  }

  TF_ASSERT_OK(session_->Extend(graph));

  const std::vector<Tensor>& input_tensors = builder.inputs();
  if (VLOG_IS_ON(1)) {
    for (const Tensor& input : input_tensors) {
      VLOG(1) << "Input: " << input.DebugString();
    }
  }

  std::vector<std::pair<string, Tensor>> expected_feeds(expected_inputs.size());
  std::vector<std::pair<string, Tensor>> test_feeds(test_inputs.size());
  ASSERT_EQ(input_tensors.size(), expected_inputs.size());
  ASSERT_EQ(input_tensors.size(), test_inputs.size());

  for (int i = 0; i < input_tensors.size(); ++i) {
    expected_feeds[i] = {expected_inputs[i], input_tensors[i]};
    test_feeds[i] = {test_inputs[i], input_tensors[i]};
  }

  std::vector<Tensor> expected_outputs, test_outputs;
  VLOG(1) << "Running expected graph";
  Status s =
      session_->Run(expected_feeds, expected_fetches, {}, &expected_outputs);
  if (!s.ok()) {
    VLOG(1) << "Expected graph failed with status: " << s << ". Skipping test";
    return;
  }
  for (const Tensor& expected : expected_outputs) {
    VLOG(1) << "Expected: " << expected.DebugString();
  }

  VLOG(1) << "Running test graph";
  TF_ASSERT_OK(session_->Run(test_feeds, test_fetches, {}, &test_outputs));

  ASSERT_EQ(expected_outputs.size(), test_outputs.size());
  for (int j = 0; s.ok() && j < test_outputs.size(); ++j) {
    s = TensorsAreClose(expected_outputs[j], test_outputs[j], atol, rtol);
  }
  TF_EXPECT_OK(s);
}

// Helper that converts 'values' to an int32 or int64 Tensor.
Tensor AsIntTensor(DataType dtype, const std::vector<int64>& values) {
  switch (dtype) {
    case DT_INT32: {
      std::vector<int32> values32(values.begin(), values.end());
      return test::AsTensor<int32>(values32);
    }
    case DT_INT64:
      return test::AsTensor<int64>(values);
    default:
      CHECK(false);
  }
}

TEST_F(OpTest, Abs) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Abs").Input(RandomTensor(type)).Attr("T", type));
  });
}

TEST_F(OpTest, Add) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Add")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, AddN) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    int n = std::uniform_int_distribution<int>(1, 5)(generator());

    auto shape = RandomDims();

    OpTestBuilder builder("AddN");
    builder.Attr("T", type);
    builder.Attr("N", n);
    for (int i = 0; i < n; ++i) {
      builder.Input(RandomTensor(type, shape));
    }
    ExpectTfAndXlaOutputsAreClose(builder);
  });
}

TEST_F(OpTest, All) {
  Repeatedly([this]() {
    Tensor data = RandomTensor(DT_BOOL);
    Tensor indices = RandomReductionIndices(data.dims());
    bool keep_dims = Choose<bool>({false, true});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("All").Input(data).Input(indices).Attr("keep_dims",
                                                             keep_dims));
  });
}

TEST_F(OpTest, Any) {
  Repeatedly([this]() {
    Tensor data = RandomTensor(DT_BOOL);
    Tensor indices = RandomReductionIndices(data.dims());
    bool keep_dims = Choose<bool>({false, true});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Any").Input(data).Input(indices).Attr("keep_dims",
                                                             keep_dims));
  });
}

TEST_F(OpTest, AvgPool) {
  Repeatedly([this]() {
    WindowedSpatialDims d = ChooseWindowedSpatialDims(2);
    std::uniform_int_distribution<int> random_int(1, 5);

    int kernel_rows = random_int(generator()),
        kernel_cols = random_int(generator());
    int stride_rows = random_int(generator()),
        stride_cols = random_int(generator());
    string padding = Choose<string>({"SAME", "VALID"});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("AvgPool")
            .Input(
                RandomTensor(DT_FLOAT, {RandomDim(1), RandomDim(kernel_rows),
                                        RandomDim(kernel_cols), RandomDim(1)}))
            .Attr("T", DT_FLOAT)
            .Attr("ksize", {1, kernel_rows, kernel_cols, 1})
            .Attr("strides", {1, stride_rows, stride_cols, 1})
            .Attr("padding", padding)
            .Attr("data_format", "NHWC"));
  });
  // TODO(phawkins): the CPU device only implements spatial pooling. Add tests
  // for batch pooling when supported.
}

TEST_F(OpTest, AvgPool3D) {
  Repeatedly([this]() {
    std::uniform_int_distribution<int> random_int(1, 5);
    std::vector<int64> input_dims, kernel_dims, stride_dims;
    for (int i = 0; i < 3; ++i) {
      kernel_dims.push_back(random_int(generator()));
      input_dims.push_back(RandomDim(kernel_dims.back()));
      stride_dims.push_back(random_int(generator()));
    }

    string padding = Choose<string>({"SAME", "VALID"});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("AvgPool3D")
            .Input(RandomTensor(DT_FLOAT, ImageDims(FORMAT_NHWC, RandomDim(1),
                                                    RandomDim(1), input_dims)))
            .Attr("T", DT_FLOAT)
            .Attr("ksize", ImageDims(FORMAT_NHWC, 1, 1, kernel_dims))
            .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, stride_dims))
            .Attr("padding", padding)
            .Attr("data_format", "NHWC"));
  });
  // TODO(phawkins): test NCHW format (not supported by CPU)
}

TEST_F(OpTest, AvgPoolGrad) {
  Repeatedly([this]() {
    int batch = RandomDim(1), features = RandomDim(1);
    WindowedSpatialDims d = ChooseWindowedSpatialDims(2);
    std::vector<int32> input_dims =
        AsInt32s(ImageDims(FORMAT_NHWC, batch, features, d.input_dims));
    std::vector<int64> output_dims =
        ImageDims(FORMAT_NHWC, batch, features, d.output_dims);
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("AvgPoolGrad")
            .Input(test::AsTensor<int32>(input_dims))
            .Input(RandomTensor(DT_FLOAT, output_dims))
            .Attr("T", DT_FLOAT)
            .Attr("ksize", ImageDims(FORMAT_NHWC, 1, 1, d.kernel_dims))
            .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims))
            .Attr("padding", d.padding == SAME ? "SAME" : "VALID")
            .Attr("data_format", "NHWC"));
  });
}

TEST_F(OpTest, AvgPool3DGrad) {
  Repeatedly([this]() {
    int batch = RandomDim(1), features = RandomDim(1);
    WindowedSpatialDims d = ChooseWindowedSpatialDims(3);
    std::vector<int32> input_dims =
        AsInt32s(ImageDims(FORMAT_NHWC, batch, features, d.input_dims));
    std::vector<int64> output_dims =
        ImageDims(FORMAT_NHWC, batch, features, d.output_dims);
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("AvgPool3DGrad")
            .Input(test::AsTensor<int32>(input_dims))
            .Input(RandomTensor(DT_FLOAT, output_dims))
            .Attr("T", DT_FLOAT)
            .Attr("ksize", ImageDims(FORMAT_NHWC, 1, 1, d.kernel_dims))
            .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims))
            .Attr("padding", d.padding == SAME ? "SAME" : "VALID")
            .Attr("data_format", "NHWC"));
  });
}

TEST_F(OpTest, BatchMatMul) {
  Repeatedly([this]() {
    std::vector<int64> output_dims = RandomDims(2, 5, 0, 7);
    int64 ndims = output_dims.size();
    int64 inner_dim = RandomDim();
    std::vector<int64> x_dims(output_dims), y_dims(output_dims);
    x_dims[ndims - 1] = inner_dim;
    y_dims[ndims - 2] = inner_dim;
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BatchMatMul")
                                      .Input(RandomTensor(DT_FLOAT, x_dims))
                                      .Input(RandomTensor(DT_FLOAT, y_dims))
                                      .Attr("T", DT_FLOAT));

    std::swap(x_dims[ndims - 1], x_dims[ndims - 2]);
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BatchMatMul")
                                      .Input(RandomTensor(DT_FLOAT, x_dims))
                                      .Input(RandomTensor(DT_FLOAT, y_dims))
                                      .Attr("T", DT_FLOAT)
                                      .Attr("adj_x", true));

    std::swap(y_dims[ndims - 1], y_dims[ndims - 2]);
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BatchMatMul")
                                      .Input(RandomTensor(DT_FLOAT, x_dims))
                                      .Input(RandomTensor(DT_FLOAT, y_dims))
                                      .Attr("T", DT_FLOAT)
                                      .Attr("adj_x", true)
                                      .Attr("adj_y", true));

    std::swap(x_dims[ndims - 1], x_dims[ndims - 2]);
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BatchMatMul")
                                      .Input(RandomTensor(DT_FLOAT, x_dims))
                                      .Input(RandomTensor(DT_FLOAT, y_dims))
                                      .Attr("T", DT_FLOAT)
                                      .Attr("adj_y", true));
  });
}

TEST_F(OpTest, BatchToSpace) {
  Repeatedly([this]() {
    const int num_block_dims = 2;
    std::vector<int64> block_dims =
        RandomDims(num_block_dims, num_block_dims, 0, 5);
    int64 block_size = RandomDim(0, 4);

    std::vector<int64> input_dims(1 + num_block_dims + 1);
    input_dims[0] = RandomDim();
    for (int i = 0; i < num_block_dims; ++i) {
      input_dims[0] *= block_size;
      input_dims[1 + i] = block_dims[i];
    }
    input_dims[1 + num_block_dims] = RandomDim();

    std::vector<int64> crop_vals;
    std::uniform_int_distribution<int> distribution(0, 4);
    for (int i = 0; i < num_block_dims; ++i) {
      // Chooses crop values; does not always choose legal values.
      crop_vals.push_back(distribution(generator()));
      crop_vals.push_back(distribution(generator()));
    }
    Tensor crops;
    CHECK(crops.CopyFrom(AsIntTensor(DT_INT32, crop_vals),
                         TensorShape({num_block_dims, 2})));

    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BatchToSpace")
                                      .Input(RandomTensor(DT_FLOAT, input_dims))
                                      .Input(crops)
                                      .Attr("T", DT_FLOAT)
                                      .Attr("block_size", block_size));
  });
}

TEST_F(OpTest, BatchToSpaceND) {
  Repeatedly([this]() {
    std::vector<int64> block_dims = RandomDims(1, 3, 0, 5);
    int num_block_dims = block_dims.size();
    std::vector<int64> remaining_dims = RandomDims(0, 3);
    std::vector<int64> block_multipliers =
        RandomDims(block_dims.size(), block_dims.size(), 0, 4);

    std::vector<int64> input_dims(1 + num_block_dims + remaining_dims.size());
    input_dims[0] = RandomDim();
    for (int i = 0; i < num_block_dims; ++i) {
      input_dims[0] *= block_dims[i];
    }
    std::copy(block_multipliers.begin(), block_multipliers.end(),
              input_dims.begin() + 1);
    std::copy(remaining_dims.begin(), remaining_dims.end(),
              input_dims.begin() + 1 + num_block_dims);

    std::vector<int64> crop_vals;
    std::uniform_int_distribution<int> distribution(0, 3);
    for (int i = 0; i < num_block_dims; ++i) {
      // Chooses crop values; does not always choose legal values.
      crop_vals.push_back(distribution(generator()));
      crop_vals.push_back(distribution(generator()));
    }
    Tensor crops;
    CHECK(crops.CopyFrom(AsIntTensor(DT_INT32, crop_vals),
                         TensorShape({num_block_dims, 2})));

    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("BatchToSpaceND")
            .Input(RandomTensor(DT_FLOAT, input_dims))
            .Input(test::AsTensor<int32>(
                std::vector<int32>(block_dims.begin(), block_dims.end())))
            .Input(crops)
            .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, BiasAdd) {
  Repeatedly([this]() {
    auto x = RandomTensor(DT_FLOAT, RandomDims(2, kDefaultMaxRank));
    auto y = RandomTensor(DT_FLOAT, {x.dim_size(x.dims() - 1)});
    // TODO(phawkins): test both data formats.
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("BiasAdd").Input(x).Input(y).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, BiasAddGrad) {
  Repeatedly([this]() {
    auto x = RandomTensor(DT_FLOAT);
    // TODO(phawkins): test both data formats.
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("BiasAddGrad").Input(x).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, BiasAddV1) {
  Repeatedly([this]() {
    auto x = RandomTensor(DT_FLOAT, RandomDims(2, kDefaultMaxRank));
    auto y = RandomTensor(DT_FLOAT, {x.dim_size(x.dims() - 1)});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("BiasAddV1").Input(x).Input(y).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, BroadcastGradientArgs) {
  Repeatedly([this]() {
    // TODO(phawkins): only int32 seems to be implemented in Tensorflow.
    // DataType type = Choose<DataType>({DT_INT32, DT_INT64});
    DataType type = DT_INT32;
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BroadcastGradientArgs")
                                      .Input(AsIntTensor(type, dims.first))
                                      .Input(AsIntTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, Cast) {
  Repeatedly([this]() {
    DataType src_type, dst_type;
    src_type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_BOOL});
    dst_type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_BOOL});
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Cast")
                                      .Input(RandomTensor(src_type))
                                      .Attr("SrcT", src_type)
                                      .Attr("DstT", dst_type));
  });
}

TEST_F(OpTest, Ceil) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Ceil")
                                      .Input(RandomTensor(DT_FLOAT))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Concat) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    int n = std::uniform_int_distribution<int>(2, 5)(generator());

    std::vector<int64> dims = RandomDims(1);
    int concat_dim =
        std::uniform_int_distribution<int32>(0, dims.size() - 1)(generator());

    OpTestBuilder builder("Concat");
    builder.Input(test::AsScalar<int32>(concat_dim));
    builder.Attr("T", type);
    builder.Attr("N", n);
    for (int i = 0; i < n; ++i) {
      std::vector<int64> shape = dims;
      shape[concat_dim] = RandomDim();
      builder.Input(RandomTensor(type, shape));
    }
    ExpectTfAndXlaOutputsAreClose(builder);
  });
}

TEST_F(OpTest, ConcatOffset) {
  Repeatedly([this]() {
    int n = std::uniform_int_distribution<int>(2, 5)(generator());

    std::vector<int64> dims = RandomDims(1);
    int concat_dim =
        std::uniform_int_distribution<int32>(0, dims.size() - 1)(generator());

    OpTestBuilder builder("ConcatOffset");
    builder.Input(test::AsScalar<int32>(concat_dim));
    builder.Attr("N", n);
    for (int i = 0; i < n; ++i) {
      std::vector<int32> shape(dims.begin(), dims.end());
      shape[concat_dim] = RandomDim();
      builder.Input(test::AsTensor<int32>(shape));
    }
    ExpectTfAndXlaOutputsAreClose(builder);
  });
}

TEST_F(OpTest, Conv2D) {
  Repeatedly([this]() {
    WindowedSpatialDims d = ChooseWindowedSpatialDims(2);
    std::uniform_int_distribution<int> random_int(1, 5);
    int features_in = random_int(generator());
    int features_out = random_int(generator());

    int64 batch = RandomDim();

    Tensor data = RandomTensor(
        DT_FLOAT, ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims));

    Tensor kernel = RandomTensor(DT_FLOAT, {d.kernel_dims[0], d.kernel_dims[1],
                                            features_in, features_out});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Conv2D")
            .Input(data)
            .Input(kernel)
            .Attr("T", DT_FLOAT)
            .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims))
            .Attr("padding", d.padding == SAME ? "SAME" : "VALID")
            .Attr("data_format", "NHWC"));
  });
}

TEST_F(OpTest, Conv2DBackpropFilter) {
  Repeatedly([this]() {
    WindowedSpatialDims d = ChooseWindowedSpatialDims(2);
    std::uniform_int_distribution<int> random_int(1, 5);
    int features_in = random_int(generator());
    int features_out = random_int(generator());
    int32 batch = RandomDim();
    Tensor activations = RandomTensor(
        DT_FLOAT, ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims));
    Tensor backprop = RandomTensor(
        DT_FLOAT, ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims));
    Tensor kernel_shape = test::AsTensor<int32>(AsInt32s(
        {d.kernel_dims[0], d.kernel_dims[1], features_in, features_out}));
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Conv2DBackpropFilter")
            .Input(activations)
            .Input(kernel_shape)
            .Input(backprop)
            .Attr("T", DT_FLOAT)
            .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims))
            .Attr("padding", d.padding == SAME ? "SAME" : "VALID")
            .Attr("data_format", "NHWC"));
  });
}

TEST_F(OpTest, Conv2DBackpropInput) {
  Repeatedly([this]() {
    WindowedSpatialDims d = ChooseWindowedSpatialDims(2);
    std::uniform_int_distribution<int> random_int(1, 5);
    int features_in = random_int(generator());
    int features_out = random_int(generator());
    int32 batch = RandomDim();
    Tensor in_shape = test::AsTensor<int32>(
        AsInt32s(ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims)));
    Tensor backprop = RandomTensor(
        DT_FLOAT, ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims));
    Tensor kernel = RandomTensor(DT_FLOAT, {d.kernel_dims[0], d.kernel_dims[1],
                                            features_in, features_out});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Conv2DBackpropInput")
            .Input(in_shape)
            .Input(kernel)
            .Input(backprop)
            .Attr("T", DT_FLOAT)
            .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims))
            .Attr("padding", d.padding == SAME ? "SAME" : "VALID")
            .Attr("data_format", "NHWC"));
  });
}

TEST_F(OpTest, Conv3D) {
  Repeatedly([this]() {
    WindowedSpatialDims d = ChooseWindowedSpatialDims(3);
    std::uniform_int_distribution<int> random_int(1, 5);
    int features_in = random_int(generator());
    int features_out = random_int(generator());
    Tensor data =
        RandomTensor(DT_FLOAT, {RandomDim(), d.input_dims[0], d.input_dims[1],
                                d.input_dims[2], features_in});

    Tensor kernel =
        RandomTensor(DT_FLOAT, {d.kernel_dims[0], d.kernel_dims[1],
                                d.kernel_dims[2], features_in, features_out});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Conv3D")
            .Input(data)
            .Input(kernel)
            .Attr("T", DT_FLOAT)
            .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims))
            .Attr("padding", d.padding == SAME ? "SAME" : "VALID"));
  });
}

TEST_F(OpTest, Conv3DBackpropFilter) {
  Repeatedly([this]() {
    WindowedSpatialDims d = ChooseWindowedSpatialDims(3);
    std::uniform_int_distribution<int> random_int(1, 5);
    int features_in = random_int(generator());
    int features_out = random_int(generator());
    int32 batch = RandomDim(1);
    Tensor activations = RandomTensor(
        DT_FLOAT, ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims));
    Tensor backprop = RandomTensor(
        DT_FLOAT, ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims));
    Tensor kernel_shape = test::AsTensor<int32>(
        AsInt32s({d.kernel_dims[0], d.kernel_dims[1], d.kernel_dims[2],
                  features_in, features_out}));
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Conv3DBackpropFilterV2")
            .Input(activations)
            .Input(kernel_shape)
            .Input(backprop)
            .Attr("T", DT_FLOAT)
            .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims))
            .Attr("padding", d.padding == SAME ? "SAME" : "VALID"));
  });
}

TEST_F(OpTest, Conv3DBackpropInput) {
  Repeatedly([this]() {
    WindowedSpatialDims d = ChooseWindowedSpatialDims(3);
    std::uniform_int_distribution<int> random_int(1, 5);
    int features_in = random_int(generator());
    int features_out = random_int(generator());
    int32 batch = RandomDim(1);
    Tensor in_shape = test::AsTensor<int32>(
        AsInt32s(ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims)));
    Tensor backprop = RandomTensor(
        DT_FLOAT, ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims));
    Tensor kernel =
        RandomTensor(DT_FLOAT, {d.kernel_dims[0], d.kernel_dims[1],
                                d.kernel_dims[2], features_in, features_out});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Conv3DBackpropInputV2")
            .Input(in_shape)
            .Input(kernel)
            .Input(backprop)
            .Attr("T", DT_FLOAT)
            .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims))
            .Attr("padding", d.padding == SAME ? "SAME" : "VALID"));
  });
}

TEST_F(OpTest, Diag) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Diag")
                                      .Input(RandomTensor(type, RandomDims(1)))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, DiagPart) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = RandomDims(1, 3);
    // Duplicate the random dims.
    std::vector<int64> doubled_dims(dims.size() * 2);
    std::copy(dims.begin(), dims.end(), doubled_dims.begin());
    std::copy(dims.begin(), dims.end(), doubled_dims.begin() + dims.size());
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("DiagPart")
                                      .Input(RandomTensor(type, doubled_dims))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, Div) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Div")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, DynamicStitch) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    int n = std::uniform_int_distribution<int>(2, 5)(generator());
    OpTestBuilder builder("DynamicStitch");
    builder.Attr("T", type);
    builder.Attr("N", n);
    std::vector<std::vector<int64>> index_dims;
    int size = 0;
    // TODO(phawkins): the XLA implementation of DynamicStitch does not
    // accept an empty set of indices.
    do {
      size = 0;
      index_dims.clear();
      for (int i = 0; i < n; ++i) {
        std::vector<int64> dims = RandomDims(0, 3, 0, 5);
        size += TensorShape(dims).num_elements();
        index_dims.push_back(dims);
      }
    } while (size == 0);

    // Shuffle the range of indices that cover the output.
    // TODO(phawkins): The documentation for DynamicStitch doesn't require that
    // the indices cover all positions of the output. The XLA implementation
    // does so require. However, the native TF implementation leaves undefined
    // values if we don't cover everything, so we can't really test that case
    // anyway.
    std::vector<int32> indices(size);
    std::iota(indices.begin(), indices.end(), 0);
    std::shuffle(indices.begin(), indices.end(), generator());

    int pos = 0;
    for (int i = 0; i < n; ++i) {
      TensorShape shape(index_dims[i]);
      Tensor t = test::AsTensor<int32>(
          gtl::ArraySlice<int32>(indices, pos, shape.num_elements()), shape);
      builder.Input(t);
      pos += t.NumElements();
    }

    std::vector<int64> constant_dims = RandomDims(0, 3, 0, 5);
    for (int i = 0; i < n; ++i) {
      std::vector<int64> dims(index_dims[i].begin(), index_dims[i].end());
      std::copy(constant_dims.begin(), constant_dims.end(),
                std::back_inserter(dims));
      Tensor t = RandomTensor(type, dims);
      builder.Input(t);
    }
    ExpectTfAndXlaOutputsAreClose(builder);
  });
}

TEST_F(OpTest, Elu) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Elu").Input(RandomTensor(DT_FLOAT)).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, EluGrad) {
  Repeatedly([this]() {
    auto dims = RandomDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("EluGrad")
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Equal) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Equal")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, Exp) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Exp").Input(RandomTensor(DT_FLOAT)).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, ExpandDims) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    Tensor in = RandomTensor(type);
    Tensor dim(DT_INT32, TensorShape());
    std::uniform_int_distribution<int32> d(-1 - in.dims(), in.dims());
    dim.scalar<int32>()() = d(generator());
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("ExpandDims").Input(in).Input(dim).Attr("T", type));
  });
}

TEST_F(OpTest, Fill) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    Tensor scalar = RandomTensor(type, {});
    std::vector<int64> dims = RandomDims();
    std::vector<int32> shape(dims.begin(), dims.end());
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Fill")
                                      .Input(test::AsTensor<int32>(shape))
                                      .Input(scalar)
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, Floor) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Floor")
                                      .Input(RandomTensor(DT_FLOAT))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, FloorDiv) {
  Repeatedly([this]() {
    DataType type = DT_INT32;
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("FloorDiv")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, FloorMod) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("FloorMod")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, Greater) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Greater")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, GreaterEqual) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("GreaterEqual")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, Reciprocal) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Reciprocal")
                                      .Input(RandomTensor(DT_FLOAT))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, L2Loss) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    // TODO(b/31644876): scalars currently crash.
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("L2Loss")
                                      .Input(RandomTensor(type, RandomDims(1)))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, Less) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Less")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, LessEqual) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("LessEqual")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, LinSpace) {
  Repeatedly([this]() {
    auto ToScalar = [](DataType type, int x) {
      if (type == DT_INT32) return test::AsScalar<int32>(x);
      return test::AsScalar<int64>(x);
    };
    std::uniform_int_distribution<int> distribution(-50, 50);
    DataType type = Choose<DataType>({DT_INT32, DT_INT64});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("LinSpace")
            .Input(RandomTensor(DT_FLOAT, {}))
            .Input(RandomTensor(DT_FLOAT, {}))
            .Input(ToScalar(type, distribution(generator())))
            .Attr("T", DT_FLOAT)
            .Attr("Tidx", type));
  });
}

TEST_F(OpTest, Log) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Log").Input(RandomTensor(DT_FLOAT)).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, LogicalAnd) {
  Repeatedly([this]() {
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("LogicalAnd")
            .Input(RandomTensor(DT_BOOL, dims.first))
            .Input(RandomTensor(DT_BOOL, dims.second)));
  });
}

TEST_F(OpTest, LogicalNot) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("LogicalNot").Input(RandomTensor(DT_BOOL)));
  });
}

TEST_F(OpTest, LogicalOr) {
  Repeatedly([this]() {
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("LogicalOr")
            .Input(RandomTensor(DT_BOOL, dims.first))
            .Input(RandomTensor(DT_BOOL, dims.second)));
  });
}

TEST_F(OpTest, LogSoftmax) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("LogSoftmax")
            .Input(RandomTensor(DT_FLOAT, RandomDims(2, 2)))
            .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, LRN) {
  Repeatedly([this]() {
    Tensor data;
    // TODO(b/31362467): Crashes with 0 dims on GPU. Re-enable when fixed.
    data = RandomTensor(DT_FLOAT, RandomDims(4, 4, 1, 8));
    // CuDNN requires depth_radius > 0.
    std::uniform_int_distribution<int> radius(1, data.dim_size(3));
    std::uniform_real_distribution<float> coeff(0.01, 2.0);
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("LRN")
                                      .Input(data)
                                      .Attr("T", DT_FLOAT)
                                      .Attr("depth_radius", radius(generator()))
                                      .Attr("bias", coeff(generator()))
                                      .Attr("alpha", coeff(generator()))
                                      .Attr("beta", coeff(generator())));
  });
}

TEST_F(OpTest, LRNGrad) {
  Repeatedly([this]() {
    // TODO(b/31362467): Crashes with 0 dims on GPU. Re-enable when fixed.
    std::vector<int64> dims = RandomDims(4, 4, 1, 8);
    Tensor input_grads = RandomTensor(DT_FLOAT, dims);
    Tensor input_image = RandomTensor(DT_FLOAT, dims);
    Tensor output_image = RandomTensor(DT_FLOAT, dims);
    // CuDNN requires depth_radius > 0.
    std::uniform_int_distribution<int> radius(1, input_grads.dim_size(3));
    std::uniform_real_distribution<float> coeff(0.0, 2.0);
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("LRNGrad")
                                      .Input(input_grads)
                                      .Input(input_image)
                                      .Input(output_image)
                                      .Attr("T", DT_FLOAT)
                                      .Attr("depth_radius", radius(generator()))
                                      .Attr("bias", coeff(generator()))
                                      .Attr("alpha", coeff(generator()))
                                      .Attr("beta", coeff(generator())));
  });
}

TEST_F(OpTest, MatMul) {
  Repeatedly([this]() {
    int64 x = RandomDim();
    int64 y = RandomDim();
    int64 z = RandomDim();

    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatMul")
                                      .Input(RandomTensor(DT_FLOAT, {x, y}))
                                      .Input(RandomTensor(DT_FLOAT, {y, z}))
                                      .Attr("T", DT_FLOAT));

    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatMul")
                                      .Input(RandomTensor(DT_FLOAT, {y, x}))
                                      .Input(RandomTensor(DT_FLOAT, {y, z}))
                                      .Attr("T", DT_FLOAT)
                                      .Attr("transpose_a", true));

    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatMul")
                                      .Input(RandomTensor(DT_FLOAT, {x, y}))
                                      .Input(RandomTensor(DT_FLOAT, {z, y}))
                                      .Attr("T", DT_FLOAT)
                                      .Attr("transpose_b", true));

    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatMul")
                                      .Input(RandomTensor(DT_FLOAT, {y, x}))
                                      .Input(RandomTensor(DT_FLOAT, {z, y}))
                                      .Attr("T", DT_FLOAT)
                                      .Attr("transpose_a", true)
                                      .Attr("transpose_b", true));
  });
}

TEST_F(OpTest, MatrixDiag) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_BOOL, DT_INT32, DT_FLOAT});
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatrixDiag")
                                      .Input(RandomTensor(type, RandomDims(1)))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, MatrixDiagPart) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_BOOL, DT_INT32, DT_FLOAT});
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatrixDiagPart")
                                      .Input(RandomTensor(type, RandomDims(2)))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, Max) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    Tensor data = RandomTensor(type);
    Tensor indices = RandomReductionIndices(data.dims());
    bool keep_dims = Choose<bool>({false, true});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Max").Input(data).Input(indices).Attr("T", type).Attr(
            "keep_dims", keep_dims));
  });
}

TEST_F(OpTest, Maximum) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Maximum")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, MaxPool) {
  Repeatedly([this]() {
    std::uniform_int_distribution<int> random_int(1, 5);
    int kernel_rows = random_int(generator()),
        kernel_cols = random_int(generator());
    int stride_rows = random_int(generator()),
        stride_cols = random_int(generator());
    string padding = Choose<string>({"SAME", "VALID"});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("MaxPool")
            .Input(
                RandomTensor(DT_FLOAT, {RandomDim(1), RandomDim(kernel_rows),
                                        RandomDim(kernel_cols), RandomDim(1)}))
            .Attr("T", DT_FLOAT)
            .Attr("ksize", {1, kernel_rows, kernel_cols, 1})
            .Attr("strides", {1, stride_rows, stride_cols, 1})
            .Attr("padding", padding)
            .Attr("data_format", "NHWC"));
  });
  // TODO(phawkins): test NCHW format (not supported by CPU)
}

TEST_F(OpTest, MaxPool3D) {
  Repeatedly([this]() {
    std::uniform_int_distribution<int> random_int(1, 5);
    std::vector<int64> input_dims;
    std::vector<int32> kernel_dims, stride_dims;
    input_dims.push_back(RandomDim(1));
    kernel_dims.push_back(1);
    stride_dims.push_back(1);
    for (int i = 0; i < 3; ++i) {
      kernel_dims.push_back(random_int(generator()));
      input_dims.push_back(RandomDim(kernel_dims.back()));
      stride_dims.push_back(random_int(generator()));
    }
    input_dims.push_back(RandomDim(1));
    kernel_dims.push_back(1);
    stride_dims.push_back(1);

    string padding = Choose<string>({"SAME", "VALID"});
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MaxPool3D")
                                      .Input(RandomTensor(DT_FLOAT, input_dims))
                                      .Attr("T", DT_FLOAT)
                                      .Attr("ksize", kernel_dims)
                                      .Attr("strides", stride_dims)
                                      .Attr("padding", padding)
                                      .Attr("data_format", "NHWC"));
  });
  // TODO(phawkins): test NCHW format (not supported by CPU)
}

TEST_F(OpTest, Mean) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    // TODO(phawkins): CPU and XLA differ output for reducing across a
    // size-0 dimension (nan vs 0). For now, require size >= 1.
    Tensor data = RandomTensor(type, RandomDims(0, kDefaultMaxRank, 1));
    Tensor indices = RandomReductionIndices(data.dims());
    bool keep_dims = Choose<bool>({false, true});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Mean").Input(data).Input(indices).Attr("T", type).Attr(
            "keep_dims", keep_dims));
  });
}

TEST_F(OpTest, Min) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    Tensor data = RandomTensor(type);
    Tensor indices = RandomReductionIndices(data.dims());
    bool keep_dims = Choose<bool>({false, true});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Min").Input(data).Input(indices).Attr("T", type).Attr(
            "keep_dims", keep_dims));
  });
}

TEST_F(OpTest, Minimum) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Minimum")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, Mod) {
  Repeatedly([this]() {
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Mod")
            .Input(RandomTensor(DT_INT32, dims.first))
            .Input(RandomTensor(DT_INT32, dims.second))
            .Attr("T", DT_INT32));
  });
}

TEST_F(OpTest, Mul) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Mul")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, Neg) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Neg").Input(RandomTensor(type)).Attr("T", type));
  });
}

TEST_F(OpTest, NotEqual) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("NotEqual")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, OneHot) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);

    std::vector<int64> dims = RandomDims();
    int num_dims = dims.size();

    int32 depth = RandomDim();

    Tensor indices(DT_INT32, TensorShape(dims));
    std::uniform_int_distribution<int32> distribution(-depth * 2, depth * 2);
    test::FillFn<int32>(&indices, [this, &distribution](int i) -> int32 {
      return distribution(generator());
    });

    int axis = std::uniform_int_distribution<int32>(-num_dims - 5,
                                                    num_dims + 5)(generator());

    OpTestBuilder builder("OneHot");
    builder.Attr("T", type);
    builder.Attr("TI", DT_INT32);
    builder.Attr("axis", axis);
    builder.Input(indices);
    builder.Input(test::AsScalar<int32>(depth));
    builder.Input(RandomTensor(type, {}));
    builder.Input(RandomTensor(type, {}));
    ExpectTfAndXlaOutputsAreClose(builder);
  });
}

TEST_F(OpTest, Pack) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    int n = std::uniform_int_distribution<int>(1, 5)(generator());

    std::vector<int64> dims = RandomDims();
    int num_dims = dims.size();
    int axis = std::uniform_int_distribution<int32>(-num_dims - 1,
                                                    num_dims)(generator());

    OpTestBuilder builder("Pack");
    builder.Attr("T", type);
    builder.Attr("N", n);
    builder.Attr("axis", axis);
    for (int i = 0; i < n; ++i) {
      builder.Input(RandomTensor(type, dims));
    }
    ExpectTfAndXlaOutputsAreClose(builder);
  });
}

// TODO(b/31741898): crashes on GPU.
TEST_F(OpTest, Pad) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    Tensor t = RandomTensor(type);

    // TODO(b/31741996): re-enable DT_INT64 when bug is fixed.
    // DataType tpaddings = Choose<DataType>({DT_INT32, DT_INT64});
    DataType tpaddings = DT_INT32;
    std::vector<int64> paddings_vec;
    std::uniform_int_distribution<int> distribution(0, 7);
    for (int i = 0; i < t.dims(); ++i) {
      paddings_vec.push_back(distribution(generator()));
      paddings_vec.push_back(distribution(generator()));
    }
    Tensor paddings;
    CHECK(paddings.CopyFrom(AsIntTensor(tpaddings, paddings_vec),
                            TensorShape({t.dims(), 2})));
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Pad").Input(t).Input(paddings).Attr("T", type).Attr(
            "Tpaddings", tpaddings));
  });
}

TEST_F(OpTest, Pow) {
  // TODO(phawkins): Feeding large DT_INT32 values to Pow() leads to
  // nontermination.
  Repeatedly([this]() {
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Pow")
            .Input(RandomTensor(DT_FLOAT, dims.first))
            .Input(RandomTensor(DT_FLOAT, dims.second))
            .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Prod) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    Tensor data = RandomTensor(type);
    Tensor indices = RandomReductionIndices(data.dims());
    bool keep_dims = Choose<bool>({false, true});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Prod").Input(data).Input(indices).Attr("T", type).Attr(
            "keep_dims", keep_dims));
  });
}

TEST_F(OpTest, Range) {
  Repeatedly([this]() {
    auto ToScalar = [](DataType type, int x) {
      if (type == DT_INT32) return test::AsScalar<int32>(x);
      if (type == DT_INT64) return test::AsScalar<int64>(x);
      if (type == DT_FLOAT) return test::AsScalar<float>(x);
      if (type == DT_DOUBLE) return test::AsScalar<double>(x);
      LOG(FATAL) << "Unknown type " << DataTypeString(type);
    };
    std::uniform_int_distribution<int> distribution(-50, 50);
    DataType tidx = Choose<DataType>({DT_INT32, DT_INT64, DT_FLOAT, DT_DOUBLE});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Range")
            .Input(ToScalar(tidx, distribution(generator())))
            .Input(ToScalar(tidx, distribution(generator())))
            .Input(ToScalar(tidx, distribution(generator())))
            .Attr("Tidx", tidx));
  });
}

TEST_F(OpTest, Rank) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Rank").Input(RandomTensor(type)).Attr("T", type));
  });
}

TEST_F(OpTest, RealDiv) {
  Repeatedly([this]() {
    DataType type = DT_FLOAT;
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("RealDiv")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, Relu) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Relu")
                                      .Input(RandomTensor(DT_FLOAT))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Relu6) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Relu6")
                                      .Input(RandomTensor(DT_FLOAT))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Relu6Grad) {
  Repeatedly([this]() {
    auto dims = RandomDims(1);
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Relu6Grad")
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, ReluGrad) {
  Repeatedly([this]() {
    auto dims = RandomDims(1);
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ReluGrad")
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Reshape) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    std::vector<int64> dims = RandomDims();
    std::bernoulli_distribution random_bool;
    std::vector<int64> dims_before, dims_after;
    for (std::vector<int64>* out : {&dims_before, &dims_after}) {
      std::shuffle(dims.begin(), dims.end(), generator());
      for (int64 dim : dims) {
        // Either add the dimension as a new dimension or merge it with the
        // previous dimension.
        if (out->empty() || random_bool(generator())) {
          out->push_back(dim);
        } else {
          out->back() *= dim;
        }
      }
    }
    Tensor data = RandomTensor(type, dims_before);
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Reshape")
            .Input(data)
            .Input(test::AsTensor<int32>(
                std::vector<int32>(dims_after.begin(), dims_after.end())))
            .Attr("T", type));
  });
}

TEST_F(OpTest, Reverse) {
  Repeatedly([this]() {
    std::vector<int64> dims = RandomDims(1);
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    int64 rank = dims.size();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Reverse")
                                      .Input(RandomTensor(type, dims))
                                      .Input(RandomTensor(DT_BOOL, {rank}))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, ReverseV2) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    Tensor data = RandomTensor(type);
    Tensor indices = RandomReductionIndices(data.dims());
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ReverseV2")
                                      .Input(data)
                                      .Input(indices)
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Round) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Round")
                                      .Input(RandomTensor(DT_FLOAT))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Rsqrt) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Rsqrt")
                                      .Input(RandomTensor(DT_FLOAT))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, RsqrtGrad) {
  Repeatedly([this]() {
    auto dims = RandomDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("RsqrtGrad")
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Shape) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Shape").Input(RandomTensor(type)).Attr("T", type));
  });
}

TEST_F(OpTest, ShapeN) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    int n = std::uniform_int_distribution<int>(1, 5)(generator());
    OpTestBuilder builder("ShapeN");
    builder.Attr("T", type);
    builder.Attr("N", n);
    for (int i = 0; i < n; ++i) {
      builder.Input(RandomTensor(type));
    }
    ExpectTfAndXlaOutputsAreClose(builder);
  });
}

TEST_F(OpTest, Sigmoid) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Sigmoid")
                                      .Input(RandomTensor(DT_FLOAT))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, SigmoidGrad) {
  Repeatedly([this]() {
    auto dims = RandomDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SigmoidGrad")
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Sign) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Sign").Input(RandomTensor(type)).Attr("T", type));
  });
}

TEST_F(OpTest, Size) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Size").Input(RandomTensor(type)).Attr("T", type));
  });
}

TEST_F(OpTest, Slice) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    Tensor data = RandomTensor(type);

    std::vector<int32> begin(data.dims()), size(data.dims());
    for (int i = 0; i < data.dims(); ++i) {
      begin[i] = std::uniform_int_distribution<int32>(
          0, data.dim_size(i))(generator());
      size[i] = std::uniform_int_distribution<int32>(
          -1, data.dim_size(i) - begin[i])(generator());
    }
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Slice")
                                      .Input(data)
                                      .Input(test::AsTensor<int32>(begin))
                                      .Input(test::AsTensor<int32>(size))
                                      .Attr("T", type)
                                      .Attr("Index", DT_INT32));
  });
}

TEST_F(OpTest, Softmax) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Softmax")
            .Input(RandomTensor(DT_FLOAT, RandomDims(2, 2)))
            .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, SoftmaxCrossEntropyWithLogits) {
  Repeatedly([this]() {
    std::vector<int64> dims = RandomDims(2, 2, 1);
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SoftmaxCrossEntropyWithLogits")
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Softplus) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Softplus")
                                      .Input(RandomTensor(DT_FLOAT))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, SoftplusGrad) {
  Repeatedly([this]() {
    std::vector<int64> dims = RandomDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SoftplusGrad")
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, SpaceToBatch) {
  Repeatedly([this]() {
    std::vector<int64> block_dims = RandomDims(4, 4, 0, 5);
    const int num_block_dims = 2;
    int64 block_size = RandomDim(0, 4);

    std::vector<int64> input_dims(1 + num_block_dims + 1);
    input_dims[0] = RandomDim();
    for (int i = 0; i < num_block_dims; ++i) {
      input_dims[1 + i] = block_dims[i] * block_size;
    }
    input_dims[1 + num_block_dims] = RandomDim();

    std::vector<int64> padding_vals;
    std::uniform_int_distribution<int> distribution(0, 7);
    for (int i = 0; i < num_block_dims; ++i) {
      int64 pad_before;
      int64 pad_after;
      do {
        pad_before = distribution(generator());
        pad_after = distribution(generator());
      } while (pad_before + pad_after > input_dims[1 + i]);
      input_dims[1 + i] -= pad_before + pad_after;
      padding_vals.push_back(pad_before);
      padding_vals.push_back(pad_after);
    }
    Tensor paddings;
    CHECK(paddings.CopyFrom(AsIntTensor(DT_INT32, padding_vals),
                            TensorShape({num_block_dims, 2})));

    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SpaceToBatch")
                                      .Input(RandomTensor(DT_FLOAT, input_dims))
                                      .Input(paddings)
                                      .Attr("T", DT_FLOAT)
                                      .Attr("block_size", block_size));
  });
}

TEST_F(OpTest, SpaceToBatchND) {
  Repeatedly([this]() {
    std::vector<int64> block_dims = RandomDims(1, 3, 0, 5);
    int num_block_dims = block_dims.size();
    std::vector<int64> remaining_dims = RandomDims(0, 3);
    std::vector<int64> block_multipliers =
        RandomDims(block_dims.size(), block_dims.size(), 0, 4);

    std::vector<int64> input_dims(1 + num_block_dims + remaining_dims.size());
    input_dims[0] = RandomDim();
    for (int i = 0; i < num_block_dims; ++i) {
      input_dims[1 + i] = block_dims[i] * block_multipliers[i];
    }
    std::copy(remaining_dims.begin(), remaining_dims.end(),
              input_dims.begin() + 1 + num_block_dims);

    std::vector<int64> padding_vals;
    std::uniform_int_distribution<int> distribution(0, 7);
    for (int i = 0; i < num_block_dims; ++i) {
      int64 pad_before;
      int64 pad_after;
      do {
        pad_before = distribution(generator());
        pad_after = distribution(generator());
      } while (pad_before + pad_after > input_dims[1 + i]);
      input_dims[1 + i] -= pad_before + pad_after;
      padding_vals.push_back(pad_before);
      padding_vals.push_back(pad_after);
    }
    Tensor paddings;
    CHECK(paddings.CopyFrom(AsIntTensor(DT_INT32, padding_vals),
                            TensorShape({num_block_dims, 2})));

    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("SpaceToBatchND")
            .Input(RandomTensor(DT_FLOAT, input_dims))
            .Input(test::AsTensor<int32>(
                std::vector<int32>(block_dims.begin(), block_dims.end())))
            .Input(paddings)
            .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, SparseMatMul) {
  Repeatedly([this]() {
    int64 x = RandomDim();
    int64 y = RandomDim();
    int64 z = RandomDim();

    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SparseMatMul")
                                      .Input(RandomTensor(DT_FLOAT, {x, y}))
                                      .Input(RandomTensor(DT_FLOAT, {y, z}))
                                      .Attr("Ta", DT_FLOAT)
                                      .Attr("Tb", DT_FLOAT));

    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SparseMatMul")
                                      .Input(RandomTensor(DT_FLOAT, {y, x}))
                                      .Input(RandomTensor(DT_FLOAT, {y, z}))
                                      .Attr("Ta", DT_FLOAT)
                                      .Attr("Tb", DT_FLOAT)
                                      .Attr("transpose_a", true));

    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SparseMatMul")
                                      .Input(RandomTensor(DT_FLOAT, {x, y}))
                                      .Input(RandomTensor(DT_FLOAT, {z, y}))
                                      .Attr("Ta", DT_FLOAT)
                                      .Attr("Tb", DT_FLOAT)
                                      .Attr("transpose_b", true));

    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SparseMatMul")
                                      .Input(RandomTensor(DT_FLOAT, {y, x}))
                                      .Input(RandomTensor(DT_FLOAT, {z, y}))
                                      .Attr("Ta", DT_FLOAT)
                                      .Attr("Tb", DT_FLOAT)
                                      .Attr("transpose_a", true)
                                      .Attr("transpose_b", true));
  });
}

TEST_F(OpTest, SparseSoftmaxCrossEntropyWithLogits) {
  Repeatedly([this]() {
    std::vector<int64> dims = RandomDims(2, 2, 1);
    int64 batch_size = dims[0];
    int64 num_classes = dims[1];

    std::vector<int32> indices(batch_size);
    for (int64 i = 0; i < batch_size; ++i) {
      indices[i] =
          std::uniform_int_distribution<int32>(0, num_classes - 1)(generator());
    }

    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("SparseSoftmaxCrossEntropyWithLogits")
            .Input(RandomTensor(DT_FLOAT, dims))
            .Input(test::AsTensor<int32>(indices))
            .Attr("T", DT_FLOAT)
            .Attr("Tlabels", DT_INT32));
  });
}

TEST_F(OpTest, Split) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    std::vector<int64> dims = RandomDims(1);
    std::uniform_int_distribution<int> ud;
    int32 dim = std::uniform_int_distribution<int32>(
        0, static_cast<int32>(dims.size()) - 1)(generator());
    int n = std::uniform_int_distribution<int>(1, 5)(generator());
    // Ensure 'dim' is evenly divisible by 'n'.
    dims[dim] /= n;
    dims[dim] *= n;
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Split")
                                      .Input(test::AsScalar<int32>(dim))
                                      .Input(RandomTensor(type, dims))
                                      .Attr("T", type)
                                      .Attr("num_split", n));
  });
}

TEST_F(OpTest, Sqrt) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Sqrt")
                                      .Input(RandomTensor(DT_FLOAT))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, SquaredDifference) {
  Repeatedly([this]() {
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("SquaredDifference")
            .Input(RandomTensor(DT_FLOAT, dims.first))
            .Input(RandomTensor(DT_FLOAT, dims.second))
            .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Square) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Square").Input(RandomTensor(type)).Attr("T", type));
  });
}

TEST_F(OpTest, Squeeze) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    Tensor t = RandomTensor(type, RandomDims(0, kDefaultMaxRank, 0, 5));
    std::bernoulli_distribution random_bool;
    std::vector<int> squeeze_dims;
    for (int i = 0; i < t.dims(); ++i) {
      if (t.dim_size(i) == 1 && random_bool(generator())) {
        squeeze_dims.push_back(i);
      }
    }
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Squeeze")
                                      .Input(t)
                                      .Attr("squeeze_dims", squeeze_dims)
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, Sub) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Sub")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, Sum) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    Tensor data = RandomTensor(type);
    Tensor indices = RandomReductionIndices(data.dims());
    bool keep_dims = Choose<bool>({false, true});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Sum").Input(data).Input(indices).Attr("T", type).Attr(
            "keep_dims", keep_dims));
  });
}

TEST_F(OpTest, StridedSlice) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    Tensor data = RandomTensor(type);

    std::vector<int32> begin(data.dims()), end(data.dims());
    std::vector<int32> strides(data.dims());
    for (int i = 0; i < data.dims(); ++i) {
      begin[i] = std::uniform_int_distribution<int32>(
          -2 * data.dim_size(i), 2 * data.dim_size(i))(generator());
      end[i] = std::uniform_int_distribution<int32>(
          -2 * data.dim_size(i), 2 * data.dim_size(i))(generator());
      // TODO(b/31360685): support strides other than 1 or -1
      strides[i] = std::bernoulli_distribution()(generator()) ? 1 : -1;
    }
    int64 max_bitmask = (1LL << data.dims()) - 1;
    std::uniform_int_distribution<int64> bitmask_distribution(0, max_bitmask);
    int64 begin_mask = bitmask_distribution(generator());
    int64 end_mask = bitmask_distribution(generator());

    // Create a ellipsis bitmask with at most one 1 bit set.
    int64 ellipsis_mask = 0;
    if (data.dims() > 0 && std::bernoulli_distribution()(generator())) {
      int ellipsis_pos =
          std::uniform_int_distribution<int>(0, data.dims() - 1)(generator());
      ellipsis_mask = 1LL << ellipsis_pos;
    }

    int64 new_axis_mask = bitmask_distribution(generator());
    int64 shrink_axis_mask = bitmask_distribution(generator());
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("StridedSlice")
            .Input(data)
            .Input(test::AsTensor<int32>(begin))
            .Input(test::AsTensor<int32>(end))
            .Input(test::AsTensor<int32>(strides))
            .Attr("T", type)
            .Attr("Index", DT_INT32)
            .Attr("begin_mask", begin_mask)
            .Attr("end_mask", end_mask)
            .Attr("ellipsis_mask", ellipsis_mask)
            .Attr("new_axis_mask", new_axis_mask)
            .Attr("shrink_axis_mask", shrink_axis_mask));
  });
}

TEST_F(OpTest, StridedSliceGrad) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);

    // Dimensions of the forward input.
    std::vector<int64> dims = RandomDims();

    std::vector<int64> begin(dims.size()), end(dims.size());
    std::vector<int64> strides(dims.size());
    for (int i = 0; i < dims.size(); ++i) {
      begin[i] = std::uniform_int_distribution<int64>(-2 * dims[i],
                                                      2 * dims[i])(generator());
      end[i] = std::uniform_int_distribution<int64>(-2 * dims[i],
                                                    2 * dims[i])(generator());
      strides[i] = std::uniform_int_distribution<int64>(
          -2 * dims[i], 2 * dims[i])(generator());
    }
    int64 max_bitmask = (1LL << dims.size()) - 1;
    std::uniform_int_distribution<int64> bitmask_distribution(0, max_bitmask);
    int64 begin_mask = bitmask_distribution(generator());
    int64 end_mask = bitmask_distribution(generator());

    // Create a ellipsis bitmask with at most one 1 bit set.
    int64 ellipsis_mask = 0;
    if (!dims.empty() && std::bernoulli_distribution()(generator())) {
      int ellipsis_pos =
          std::uniform_int_distribution<int>(0, dims.size() - 1)(generator());
      ellipsis_mask = 1LL << ellipsis_pos;
    }

    int64 new_axis_mask = bitmask_distribution(generator());
    int64 shrink_axis_mask = bitmask_distribution(generator());

    // TODO(phawkins): use shape inference for the forward op to compute the
    // gradient shape for the backward op. At present, there is a low
    // probability of the golden op succeeding.
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("StridedSliceGrad")
            .Input(test::AsTensor<int64>(dims))
            .Input(test::AsTensor<int64>(begin))
            .Input(test::AsTensor<int64>(end))
            .Input(test::AsTensor<int64>(strides))
            .Input(RandomTensor(type, RandomDims(1)))
            .Attr("T", type)
            .Attr("Index", DT_INT64)
            .Attr("begin_mask", begin_mask)
            .Attr("end_mask", end_mask)
            .Attr("ellipsis_mask", ellipsis_mask)
            .Attr("new_axis_mask", new_axis_mask)
            .Attr("shrink_axis_mask", shrink_axis_mask));
  });
}

TEST_F(OpTest, Tanh) {
  Repeatedly([this]() {
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Tanh")
                                      .Input(RandomTensor(DT_FLOAT))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, TanhGrad) {
  Repeatedly([this]() {
    auto dims = RandomDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("TanhGrad")
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Input(RandomTensor(DT_FLOAT, dims))
                                      .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Tile) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    Tensor t = RandomTensor(type, RandomDims(1));
    std::vector<int32> multiples(t.dims());
    for (int i = 0; i < t.dims(); ++i) {
      multiples[i] = std::uniform_int_distribution<int>(1, 3)(generator());
    }
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Tile")
                                      .Input(t)
                                      .Input(test::AsTensor<int32>(multiples))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, Transpose) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>(kAllXlaTypes);
    Tensor data = RandomTensor(type);
    std::vector<int32> perm(data.dims());
    std::iota(perm.begin(), perm.end(), 0);
    std::shuffle(perm.begin(), perm.end(), generator());
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Transpose")
                                      .Input(data)
                                      .Input(test::AsTensor<int32>(perm))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, TruncateDiv) {
  Repeatedly([this]() {
    DataType type = DT_INT32;
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("TruncateDiv")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, TruncateMod) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    auto dims = BroadcastableDims();
    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("TruncateMod")
                                      .Input(RandomTensor(type, dims.first))
                                      .Input(RandomTensor(type, dims.second))
                                      .Attr("T", type));
  });
}

TEST_F(OpTest, ZerosLike) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("ZerosLike").Input(RandomTensor(type)).Attr("T", type));
  });
}

TEST_F(OpTest, OnesLike) {
  Repeatedly([this]() {
    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
    ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("OnesLike").Input(RandomTensor(type)).Attr("T", type));
  });
}

}  // anonymous namespace
}  // namespace tensorflow

int main(int argc, char** argv) {
  tensorflow::tf_xla_test_device_ptr = new tensorflow::string("GPU:0");
  std::vector<tensorflow::Flag> flag_list = {
      tensorflow::Flag(
          "tf_xla_random_seed", &tensorflow::tf_xla_random_seed,
          "Random seed to use for XLA tests. <= 0 means choose a seed "
          "nondetermistically."),
      // TODO(phawkins): it might make more sense to run each test up to a
      // configurable time bound.
      tensorflow::Flag("tf_xla_test_repetitions",
                       &tensorflow::tf_xla_test_repetitions,
                       "Number of repetitions for each test."),
      tensorflow::Flag("tf_xla_test_device", tensorflow::tf_xla_test_device_ptr,
                       "Tensorflow device type to use for test"),
      tensorflow::Flag("tf_xla_test_use_jit", &tensorflow::tf_xla_test_use_jit,
                       "Use JIT compilation for the operator under test"),
  };
  tensorflow::string usage = tensorflow::Flags::Usage(argv[0], flag_list);
  const bool parse_result = tensorflow::Flags::Parse(&argc, argv, flag_list);
  if (!parse_result) {
    LOG(ERROR) << "\n" << usage;
    return 2;
  }
  testing::InitGoogleTest(&argc, argv);
  if (argc > 1) {
    LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage;
    return 2;
  }
  // XLA devices register kernels at construction time; create all known devices
  // to make sure the kernels are registered.
  std::vector<tensorflow::Device*> devices;
  TF_CHECK_OK(tensorflow::DeviceFactory::AddDevices(
      tensorflow::SessionOptions(), "", &devices));
  tensorflow::DeviceMgr device_mgr(devices);

  tensorflow::Device* ignored;
  TF_QCHECK_OK(
      device_mgr.LookupDevice(*tensorflow::tf_xla_test_device_ptr, &ignored))
      << "Unknown test device (" << *tensorflow::tf_xla_test_device_ptr
      << "). Did you build in the right configuration (e.g., is CUDA enabled)?";

  return RUN_ALL_TESTS();
}
