/* 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:
// * 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 <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <initializer_list>
#include <iterator>
#include <limits>
#include <memory>
#include <numeric>
#include <optional>
#include <random>
#include <string>
#include <utility>
#include <vector>

#include "absl/algorithm/container.h"
#include "absl/container/fixed_array.h"
#include "absl/container/flat_hash_set.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/jit/defs.h"
#include "tensorflow/compiler/jit/flags.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/status.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/framework/device.h"
#include "tensorflow/core/framework/device_factory.h"
#include "tensorflow/core/framework/kernel_shape_util.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/numeric_types.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/bfloat16.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/config.pb.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/padding.h"
#include "tensorflow/core/util/tensor_format.h"

namespace tensorflow {
namespace {

// Command line flags: see main() below.
int64_t tf_xla_random_seed = 0;
int32_t tf_xla_test_repetitions = 20;
int64_t tf_xla_max_tensor_size = 10000LL;
std::string* tf_xla_test_device_ptr;       // initial value set in main()
std::string* tf_xla_reference_device_ptr;  // initial value set in main()
bool tf_xla_test_use_jit = true;
bool tf_xla_test_use_mlir = false;

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

constexpr std::array<DataType, 5> kAllXlaTypes = {
    {DT_INT32, DT_INT64, DT_FLOAT, DT_BOOL, DT_COMPLEX64}};
constexpr std::array<DataType, 4> kAllNumberTypes = {
    {DT_INT32, DT_INT64, DT_FLOAT, DT_COMPLEX64}};

// 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 std::string& op_name);

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

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

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

  // Add variadic input tensors as Placehodler nodes.
  OpTestBuilder& VariadicInput(const std::vector<Tensor>& tensor);

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

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

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

  struct InputDescription {
    Tensor tensor;

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

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

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

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

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

OpTestBuilder& OpTestBuilder::RandomInput(DataType type) {
  VLOG(1) << "Adding random input: " << type;
  InputDescription input;
  input.type = type;
  inputs_.push_back(input);
  return *this;
}

OpTestBuilder& OpTestBuilder::RandomInput(DataType type,
                                          std::vector<int64_t> dims) {
  VLOG(1) << "Adding input: " << type << " " << TensorShape(dims).DebugString();
  InputDescription input;
  input.type = type;
  input.has_dims = true;
  input.dims = std::move(dims);
  inputs_.push_back(input);
  return *this;
}

OpTestBuilder& OpTestBuilder::RandomUniqueInput(DataType type,
                                                std::vector<int64_t> dims) {
  VLOG(1) << "Adding input: " << type << " " << TensorShape(dims).DebugString();
  InputDescription input;
  input.type = type;
  input.has_dims = true;
  input.needs_unique_values = true;
  input.dims = std::move(dims);
  inputs_.push_back(input);
  return *this;
}

OpTestBuilder& OpTestBuilder::VariadicInput(
    const std::vector<Tensor>& tensors) {
  VLOG(1) << "Adding variadic input of length " << tensors.size() << ":";
  for (auto& t : tensors) {
    Input(t);
  }
  return *this;
}

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

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

absl::Status OpTestBuilder::BuildGraph(
    const std::string& name_prefix, const std::string& device, bool use_jit,
    GraphDef* graphdef, NodeDef** test_node_def,
    std::vector<std::string>* inputs, std::vector<std::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(absl::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();
    std::string name = absl::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();
    std::string name = absl::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 absl::OkStatus();
}

// 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();

  enum TestResult {
    // The test saw an unrecoverable error. Don't try any more runs.
    kFatalError,
    // The parameters of the test were invalid (e.g., the "golden"
    // implementation failed, or the parameters are oversize). Reruns are ok.
    kInvalid,
    // The test ran successfully, and we have a verdict. Does *not* mean the
    // test passed.
    kOk,
  };

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

  // Select a random element from 'candidates'.
  template <typename T>
  T Choose(absl::Span<const T> candidates);

  static constexpr int kDefaultMaxRank = 5;
  static constexpr int64_t kDefaultMaxDimensionSize = 256LL;

  // Returns true if 'dims' have a size less than tf_xla_max_tensor_size.
  bool TensorSizeIsOk(absl::Span<const int64_t> dims);

  // Returns a random dimension size, in the range [min, max).
  int64_t RandomDim(int64_t min = 0, int64_t 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_t> RandomDims(int min_rank = 0,
                                  int max_rank = kDefaultMaxRank,
                                  int64_t min_size = 0,
                                  int64_t max_size = kDefaultMaxDimensionSize);

  // Given a shape 'dims', build dimensions that are broadcastable to 'dims'.
  std::vector<int64_t> BroadcastableToDims(std::vector<int64_t> dims);

  // Given a shape 'dims', build a pair of dimensions such that one broadcasts
  // to the other.
  std::pair<std::vector<int64_t>, std::vector<int64_t>> BroadcastableDims(
      std::vector<int64_t> 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_t>, std::vector<int64_t>> 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, bool needs_unique_values,
                      absl::Span<const int64_t> shape);
  Tensor RandomTensor(DataType dtype);

  // Like RandomTensor, but uses values >= 0.
  Tensor RandomNonNegativeTensor(DataType dtype,
                                 absl::Span<const int64_t> shape);
  Tensor RandomNonNegativeTensor(DataType dtype);

  // Like RandomTensor, but all values are in the range [lo, hi].
  template <typename T>
  Tensor RandomBoundedTensor(DataType dtype, T lo, T hi,
                             bool needs_unique_values,
                             absl::Span<const int64_t> shape);
  template <typename T>
  Tensor RandomBoundedTensor(DataType dtype, T lo, T hi,
                             bool needs_unique_values);

  // Like RandomTensor, but the value at index i is in the range [lo[i], hi[i]].
  Tensor RandomBoundedTensor(DataType dtype, Tensor lo, Tensor hi);

  // Like RandomTensor, but return a pair {left, right} with
  // left[i] <= right[i].
  std::pair<Tensor, Tensor> RandomLteTensors(DataType dtype,
                                             absl::Span<const int64_t> shape);
  std::pair<Tensor, Tensor> RandomLteTensors(DataType dtype);

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

  // Returns a random bit.
  bool RandomBool();

  // Randomly choose a seed for a random number generator.
  int64_t RandomSeed();

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

  struct BatchMatMulArguments {
    std::vector<int64_t> lhs_dims;
    std::vector<int64_t> rhs_dims;
    DataType dtype;
    bool adj_lhs;
    bool adj_rhs;
  };
  // Choose arguments for the tf.BatchMatMul{V2} ops.
  BatchMatMulArguments ChooseBatchMatMulArguments(bool broadcastable_batch);

  struct ConcatArguments {
    std::vector<Tensor> values;
    Tensor axis;
    int n;
    DataType type;
    DataType type_idx;
  };
  // Choose arguments for the tf.Concat{V2} ops.
  ConcatArguments ChooseConcatArguments(bool int64_idx_allowed);

  struct EinsumArguments {
    std::vector<int64_t> lhs_dims;
    std::vector<int64_t> rhs_dims;
    DataType type;
    std::string equation;
  };
  // Choose arguments for the tf.{Xla}Einsum ops.
  EinsumArguments ChooseEinsumArguments();

  struct GatherArguments {
    int64_t batch_dims;
    DataType axis_type;
    DataType indices_type;
    DataType params_type;
    std::vector<int64_t> params_shape;
    Tensor indices;
    Tensor axis;
  };
  // Choose arguments for the tf.Gather{V2} ops.
  GatherArguments ChooseGatherArguments(bool axis_0);

  struct PadArguments {
    DataType input_type;
    DataType paddings_type;
    std::vector<int64_t> input_shape;
    Tensor paddings;
    Tensor constant_values;
  };
  // Choose arguments for the tf.Pad{V2} ops.
  PadArguments ChoosePadArguments();

  struct ScatterArguments {
    DataType type;
    DataType indices_type;
    Tensor indices;
    Tensor updates;
    std::vector<int64_t> shape;
  };
  // Choose arguments for ScatterNd and TensorScatterUpdate.
  ScatterArguments ChooseScatterArguments();

  struct SliceArguments {
    DataType type;
    DataType indices_type;
    std::vector<int64_t> shape;
    Tensor indices;
    std::vector<int64_t> size;
  };
  // Choose arguments for the tf.{XlaDynamicUpdate}Slice ops.
  SliceArguments ChooseSliceArguments(bool neg_one_size);

  struct XlaDotArguments {
    std::vector<int64_t> lhs_dims;
    std::vector<int64_t> rhs_dims;
    std::string dnums_encoded;
    std::string precision_config_encoded;
    DataType dtype;
  };
  // Choose arguments for tf.XlaDot operation.
  XlaDotArguments ChooseXlaDotArguments();

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

  // Converts an int64 vector to an int32 vector.
  std::vector<int32_t> AsInt32s(const std::vector<int64_t>& 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.
  TestResult 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_t 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(ERROR) << "Random seed for test case: " << seed
             << ". To reproduce the "
                "results of this test, pass flag --tf_xla_random_seed="
             << seed;
  generator_ = std::make_unique<std::mt19937>(seed);
}

namespace {
template <typename T>
Tensor TensorFromValues(DataType dtype, absl::Span<const int64_t> shape,
                        absl::Span<T> vals) {
  Tensor tensor(dtype, TensorShape(shape));
  test::FillValues<T>(&tensor, vals);
  return tensor;
}

int64_t ShapeNumVals(absl::Span<const int64_t> shape) {
  int64_t num_vals = 1;
  for (int i = 0; i < shape.size(); ++i) {
    num_vals *= shape[i];
  }
  return num_vals;
}
}  // namespace

// TensorGenerator is an abstract class that has one implementing class for each
// (DataType,T) pair. The implementing class implements RandomVals, which is
// the only Tensor generation code that is specific to the DataType.
template <typename T>
class TensorGenerator {
 public:
  explicit TensorGenerator(OpTest& test) : test_(test) {}
  virtual ~TensorGenerator() = default;
  virtual DataType dtype() = 0;
  virtual void RandomVals(std::optional<T> lo, std::optional<T> hi,
                          bool needs_unique_values,
                          absl::FixedArray<T>& vals) = 0;

  Tensor RandomTensor(std::optional<T> lo, std::optional<T> hi,
                      bool needs_unique_values,
                      absl::Span<const int64_t> shape) {
    absl::FixedArray<T> vals(ShapeNumVals(shape));
    RandomVals(lo, hi, needs_unique_values, vals);
    return TensorFromValues<T>(dtype(), shape, absl::Span<T>(vals));
  }

  std::pair<Tensor, Tensor> RandomLteTensors(absl::Span<const int64_t> shape) {
    int64_t num_vals = ShapeNumVals(shape);
    absl::FixedArray<T> less(num_vals);
    RandomVals({}, {}, false, less);
    absl::FixedArray<T> greater(num_vals);
    RandomVals({}, {}, false, greater);
    for (int i = 0; i < num_vals; ++i) {
      if (less[i] > greater[i]) {
        std::swap(less[i], greater[i]);
      }
    }
    std::pair<Tensor, Tensor> pair(
        TensorFromValues<T>(dtype(), shape, absl::Span<T>(less)),
        TensorFromValues<T>(dtype(), shape, absl::Span<T>(greater)));
    return pair;
  }

 protected:
  OpTest& test_;
};

class TensorGeneratorFloat : public TensorGenerator<float> {
 public:
  explicit TensorGeneratorFloat(OpTest& test) : TensorGenerator(test) {}
  DataType dtype() override { return DT_FLOAT; }
  void RandomVals(std::optional<float> lo, std::optional<float> hi,
                  bool needs_unique_values,
                  absl::FixedArray<float>& vals) override {
    absl::flat_hash_set<float> already_generated;
    std::uniform_real_distribution<float> distribution(lo.value_or(-1.0f),
                                                       hi.value_or(1.0f));
    for (int64_t i = 0; i < vals.size(); ++i) {
      float generated;
      do {
        generated = distribution(test_.generator());
      } while (needs_unique_values &&
               !already_generated.insert(generated).second);
      vals[i] = (generated);
    }
  }
};

class TensorGeneratorDouble : public TensorGenerator<double> {
 public:
  explicit TensorGeneratorDouble(OpTest& test) : TensorGenerator(test) {}
  DataType dtype() override { return DT_DOUBLE; }
  void RandomVals(std::optional<double> lo, std::optional<double> hi,
                  bool needs_unique_values,
                  absl::FixedArray<double>& vals) override {
    absl::flat_hash_set<double> already_generated;
    std::uniform_real_distribution<double> distribution(lo.value_or(-1.0),
                                                        hi.value_or(1.0));
    for (int64_t i = 0; i < vals.size(); ++i) {
      double generated;
      do {
        generated = distribution(test_.generator());
      } while (needs_unique_values &&
               !already_generated.insert(generated).second);
      vals[i] = generated;
    }
  }
};

class TensorGeneratorComplex64 : public TensorGenerator<complex64> {
 public:
  explicit TensorGeneratorComplex64(OpTest& test) : TensorGenerator(test) {}
  DataType dtype() override { return DT_COMPLEX64; }
  void RandomVals(std::optional<complex64> lo, std::optional<complex64> hi,
                  bool needs_unique_values,
                  absl::FixedArray<complex64>& vals) override {
    absl::flat_hash_set<std::pair<float, float>> already_generated;
    if (lo || hi) {
      LOG(FATAL) << "Lower or upper bounds are not supported for complex64.";
    }
    std::uniform_real_distribution<float> distribution(-1.0f, 1.0f);
    for (int64_t i = 0; i < vals.size(); ++i) {
      complex64 generated;
      do {
        generated = complex64(distribution(test_.generator()),
                              distribution(test_.generator()));
      } while (needs_unique_values &&
               !already_generated
                    .insert(std::make_pair(generated.real(), generated.imag()))
                    .second);
      vals[i] = generated;
    }
  }
};

class TensorGeneratorInt32 : public TensorGenerator<int32_t> {
 public:
  explicit TensorGeneratorInt32(OpTest& test) : TensorGenerator(test) {}
  DataType dtype() override { return DT_INT32; }
  void RandomVals(std::optional<int32_t> lo, std::optional<int32_t> hi,
                  bool needs_unique_values,
                  absl::FixedArray<int32_t>& vals) override {
    absl::flat_hash_set<int32_t> already_generated;
    std::uniform_int_distribution<int32_t> distribution(lo.value_or(-(1 << 20)),
                                                        hi.value_or(1 << 20));
    for (int64_t i = 0; i < vals.size(); ++i) {
      int32_t generated;
      do {
        generated = distribution(test_.generator());
      } while (needs_unique_values &&
               !already_generated.insert(generated).second);
      vals[i] = generated;
    }
  }
};

class TensorGeneratorInt64 : public TensorGenerator<int64_t> {
 public:
  explicit TensorGeneratorInt64(OpTest& test) : TensorGenerator(test) {}
  DataType dtype() override { return DT_INT64; }
  void RandomVals(std::optional<int64_t> lo, std::optional<int64_t> hi,
                  bool needs_unique_values,
                  absl::FixedArray<int64_t>& vals) override {
    absl::flat_hash_set<int64_t> already_generated;
    std::uniform_int_distribution<int64_t> distribution(
        lo.value_or(-(1LL << 40)), hi.value_or(1LL << 40));
    for (int64_t i = 0; i < vals.size(); ++i) {
      int64_t generated;
      do {
        generated = distribution(test_.generator());
      } while (needs_unique_values &&
               !already_generated.insert(generated).second);
      vals[i] = generated;
    }
  }
};

class TensorGeneratorBool : public TensorGenerator<bool> {
 public:
  explicit TensorGeneratorBool(OpTest& test) : TensorGenerator(test) {}
  DataType dtype() override { return DT_BOOL; }
  void RandomVals(std::optional<bool> lo, std::optional<bool> hi,
                  bool needs_unique_values,
                  absl::FixedArray<bool>& vals) override {
    absl::flat_hash_set<bool> already_generated;
    if (lo || hi) {
      LOG(FATAL) << "Lower or upper bounds are not supported for bool.";
    }
    std::bernoulli_distribution distribution;
    for (int64_t i = 0; i < vals.size(); ++i) {
      bool generated;
      do {
        generated = distribution(test_.generator());
      } while (needs_unique_values &&
               !already_generated.insert(generated).second);
      vals[i] = generated;
    }
  }
};

void OpTest::Repeatedly(const std::function<TestResult()>& fn) {
  int const max_repetitions = tf_xla_test_repetitions;
  int valid_test_runs = 0;
  // We run up to 100 * max_repetitions times; the idea is that if we roll the
  // dice enough times we will find some valid parameters. We want to put an
  // upper limit on the number iterations just in case the probability of
  // finding feasible parameters is very low.
  for (int i = 0; !HasFailure() && i < max_repetitions * 100 &&
                  valid_test_runs < max_repetitions;
       ++i) {
    TestResult result = fn();
    switch (result) {
      case kOk:
        ++valid_test_runs;
        break;

      case kFatalError:
        ASSERT_TRUE(false) << "Test had fatal failure";
        return;

      case kInvalid:
        break;
    }
  }
  if (!HasFailure()) {
    EXPECT_GE(valid_test_runs, max_repetitions)
        << "Not enough test instances passed; this means that either the "
           "golden implementation is buggy or the operator harness is not "
           "producing well-formed test cases with a high probability.";
  }
}

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

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

bool OpTest::TensorSizeIsOk(absl::Span<const int64_t> dims) {
  int64_t size = 1LL;
  for (int64_t dim : dims) {
    size *= dim;
  }
  return size < tf_xla_max_tensor_size;
}

std::vector<int64_t> OpTest::RandomDims(int min_rank, int max_rank,
                                        int64_t min_size, int64_t 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_t> dims(rank);
  if (rank == 0) {
    return dims;
  }
  int64_t per_dim_limit = std::pow(tf_xla_max_tensor_size, 1.0 / rank);
  int64_t per_dim_max = std::min(max_size, per_dim_limit);
  std::generate(dims.begin(), dims.end(), [this, min_size, per_dim_max]() {
    return RandomDim(min_size, per_dim_max);
  });
  CHECK(TensorSizeIsOk(dims));  // Crash OK
  return dims;
}

bool OpTest::RandomBool() {
  std::bernoulli_distribution d(0.5);
  return d(generator());
}

int64_t OpTest::RandomSeed() {
  std::uniform_int_distribution<int64_t> seed_dist(
      std::numeric_limits<int64_t>::min(), std::numeric_limits<int64_t>::max());
  int64_t seed = seed_dist(generator());
  if (seed == 0) return 1;
  return seed;
}

Tensor OpTest::RandomTensor(DataType dtype, bool needs_unique_values,
                            absl::Span<const int64_t> shape) {
  switch (dtype) {
    case DT_FLOAT:
      return TensorGeneratorFloat(*this).RandomTensor(
          {}, {}, needs_unique_values, shape);
    case DT_DOUBLE:
      return TensorGeneratorDouble(*this).RandomTensor(
          {}, {}, needs_unique_values, shape);
    case DT_COMPLEX64:
      return TensorGeneratorComplex64(*this).RandomTensor(
          {}, {}, needs_unique_values, shape);
    case DT_INT32:
      return TensorGeneratorInt32(*this).RandomTensor(
          {}, {}, needs_unique_values, shape);
    case DT_INT64:
      return TensorGeneratorInt64(*this).RandomTensor(
          {}, {}, needs_unique_values, shape);
    case DT_BOOL:
      return TensorGeneratorBool(*this).RandomTensor(
          {}, {}, needs_unique_values, shape);
    default:
      LOG(FATAL) << "Unimplemented type " << dtype << " in RandomTensor";
  }
}

Tensor OpTest::RandomTensor(DataType dtype) {
  return RandomTensor(dtype, /*needs_unique_values=*/false, RandomDims());
}

Tensor OpTest::RandomNonNegativeTensor(DataType dtype,
                                       absl::Span<const int64_t> shape) {
  switch (dtype) {
    case DT_FLOAT:
      return TensorGeneratorFloat(*this).RandomTensor({0.0f}, {}, false, shape);
    case DT_DOUBLE:
      return TensorGeneratorDouble(*this).RandomTensor({0.0}, {}, false, shape);
    case DT_INT32:
      return TensorGeneratorInt32(*this).RandomTensor({0}, {}, false, shape);
    case DT_INT64:
      return TensorGeneratorInt64(*this).RandomTensor({0}, {}, false, shape);
    default:
      LOG(FATAL) << "Unimplemented type " << dtype
                 << " in RandomNonNegativeTensor";
  }
}

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

template <typename T>
Tensor OpTest::RandomBoundedTensor(DataType dtype, T lo, T hi,
                                   bool needs_unique_values,
                                   absl::Span<const int64_t> shape) {
  switch (dtype) {
    case DT_FLOAT:
      return TensorGeneratorFloat(*this).RandomTensor(
          {lo}, {hi}, needs_unique_values, shape);
    case DT_DOUBLE:
      return TensorGeneratorDouble(*this).RandomTensor(
          {lo}, {hi}, needs_unique_values, shape);
    case DT_INT32:
      return TensorGeneratorInt32(*this).RandomTensor(
          {lo}, {hi}, needs_unique_values, shape);
    case DT_INT64:
      return TensorGeneratorInt64(*this).RandomTensor(
          {lo}, {hi}, needs_unique_values, shape);
    default:
      LOG(FATAL) << "RandomBoundedTensor does not support type " << dtype
                 << ".";
  }
}

template <typename T>
Tensor OpTest::RandomBoundedTensor(DataType dtype, T lo, T hi,
                                   bool needs_unique_values) {
  return RandomBoundedTensor<T>(dtype, lo, hi, needs_unique_values,
                                RandomDims());
}

Tensor OpTest::RandomBoundedTensor(DataType dtype, Tensor lo, Tensor hi) {
  TensorShape shape = lo.shape();
  if (hi.shape() != shape) {
    LOG(FATAL) << "hi and lo do not have the same shape in RandomBoundedTensor";
  }
  if (hi.dtype() != dtype) {
    LOG(FATAL) << "hi does not have the expected dtype in RandomBoundedTensor";
  }
  if (lo.dtype() != dtype) {
    LOG(FATAL) << "lo does not have the expected dtype in RandomBoundedTensor";
  }
  Tensor tensor(dtype, shape);
  switch (dtype) {
    case DT_FLOAT: {
      auto lo_flat = lo.flat<float>();
      auto hi_flat = hi.flat<float>();
      test::FillFn<float>(&tensor, [this, &lo_flat, &hi_flat](int i) -> float {
        std::uniform_real_distribution<float> distribution(lo_flat(i),
                                                           hi_flat(i));
        return distribution(generator());
      });
      break;
    }
    case DT_DOUBLE: {
      auto lo_flat = lo.flat<double>();
      auto hi_flat = hi.flat<double>();
      test::FillFn<double>(
          &tensor, [this, &lo_flat, &hi_flat](int i) -> double {
            std::uniform_real_distribution<double> distribution(lo_flat(i),
                                                                hi_flat(i));
            return distribution(generator());
          });
      break;
    }
    case DT_INT32: {
      auto lo_flat = lo.flat<int32_t>();
      auto hi_flat = hi.flat<int32_t>();
      test::FillFn<int32_t>(
          &tensor, [this, &lo_flat, &hi_flat](int i) -> int32_t {
            std::uniform_int_distribution<int32_t> distribution(lo_flat(i),
                                                                hi_flat(i));
            return distribution(generator());
          });
      break;
    }
    case DT_INT64: {
      auto lo_flat = lo.flat<int64_t>();
      auto hi_flat = hi.flat<int64_t>();
      test::FillFn<int64_t>(
          &tensor, [this, &lo_flat, &hi_flat](int i) -> int64_t {
            std::uniform_int_distribution<int64_t> distribution(lo_flat(i),
                                                                hi_flat(i));
            return distribution(generator());
          });
      break;
    }
    default:
      LOG(FATAL) << "RandomBoundedTensor does not support type " << dtype
                 << ".";
  }
  return tensor;
}

std::pair<Tensor, Tensor> OpTest::RandomLteTensors(
    DataType dtype, absl::Span<const int64_t> shape) {
  switch (dtype) {
    case DT_FLOAT:
      return TensorGeneratorFloat(*this).RandomLteTensors(shape);
    case DT_DOUBLE:
      return TensorGeneratorDouble(*this).RandomLteTensors(shape);
    case DT_COMPLEX64:
      LOG(FATAL) << "RandomLteTensors unavailable for DT_COMPLEX64";
      break;
    case DT_INT32:
      return TensorGeneratorInt32(*this).RandomLteTensors(shape);
    case DT_INT64:
      return TensorGeneratorInt64(*this).RandomLteTensors(shape);
    case DT_BOOL:
      LOG(FATAL) << "RandomLteTensors unavailable for DT_BOOL";
      break;
    default:
      LOG(FATAL) << "Unimplemented type " << dtype << " in RandomLteTensors";
  }
  Tensor tensor(dtype, TensorShape(shape));
  return std::pair<Tensor, Tensor>(tensor, tensor);
}

std::pair<Tensor, Tensor> OpTest::RandomLteTensors(DataType dtype) {
  return RandomLteTensors(dtype, RandomDims());
}

std::vector<int64_t> OpTest::BroadcastableToDims(std::vector<int64_t> dims) {
  if (dims.empty()) return 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_t> bdims(dims.begin() + skip, dims.end());

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

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

std::pair<std::vector<int64_t>, std::vector<int64_t>> OpTest::BroadcastableDims(
    std::vector<int64_t> dims) {
  auto bdims = BroadcastableToDims(dims);
  // Possibly swap the roles of 'dims' and 'bdims'.
  std::bernoulli_distribution random_bool;
  if (random_bool(generator())) {
    dims.swap(bdims);
  }
  return {dims, bdims};
}

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

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

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

OpTest::BatchMatMulArguments OpTest::ChooseBatchMatMulArguments(
    bool broadcastable_batch) {
  BatchMatMulArguments a;
  a.dtype = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});

  int64_t min_size = 0;
  int64_t max_size = 7;
  auto batch_dims_to = RandomDims(0, 3, min_size, max_size);
  int rank = batch_dims_to.size() + 2;
  std::pair<std::vector<int64_t>, std::vector<int64_t>> batch_dims_nobcast(
      batch_dims_to, batch_dims_to);
  auto batch_dims = broadcastable_batch ? BroadcastableDims(batch_dims_to)
                                        : batch_dims_nobcast;
  std::vector<int64_t> lhs_dims(batch_dims.first), rhs_dims(batch_dims.second);
  int64_t inner_dim = RandomDim();
  lhs_dims.push_back(RandomDim(min_size, max_size));
  lhs_dims.push_back(inner_dim);
  rhs_dims.push_back(inner_dim);
  rhs_dims.push_back(RandomDim(min_size, max_size));

  std::bernoulli_distribution random_bool;
  a.adj_lhs = random_bool(generator());
  a.adj_rhs = random_bool(generator());
  if (a.adj_lhs) {
    std::swap(lhs_dims[rank - 1], lhs_dims[rank - 2]);
  }
  if (a.adj_rhs) {
    std::swap(rhs_dims[rank - 1], rhs_dims[rank - 2]);
  }

  a.lhs_dims = lhs_dims;
  a.rhs_dims = rhs_dims;
  return a;
}

OpTest::ConcatArguments OpTest::ChooseConcatArguments(bool int64_idx_allowed) {
  ConcatArguments a;

  std::bernoulli_distribution random_bool;
  bool use_int64_idx = random_bool(generator());

  a.type = Choose<DataType>(kAllXlaTypes);
  a.type_idx = use_int64_idx ? DT_INT64 : DT_INT32;
  a.n = std::uniform_int_distribution<int>(2, 4)(generator());

  std::vector<int64_t> dims = RandomDims(1, 4, 0, 64);

  int axis =
      std::uniform_int_distribution<int32_t>(0, dims.size() - 1)(generator());
  a.axis = use_int64_idx ? test::AsScalar<int64_t>(axis)
                         : test::AsScalar<int32_t>(axis);

  for (int i = 0; i < a.n; ++i) {
    std::vector<int64_t> shape = dims;
    shape[axis] = RandomDim(0, 64);
    a.values.push_back(RandomTensor(a.type, false, shape));
  }

  return a;
}

OpTest::EinsumArguments OpTest::ChooseEinsumArguments() {
  EinsumArguments a;

  enum EinsumType { matmul, batchmatmul, dot, outer };
  int op_kind = Choose<int>({matmul, batchmatmul, dot, outer});
  switch (op_kind) {
    case matmul:
    case batchmatmul: {
      std::vector<int64_t> dims;
      if (op_kind == matmul) {
        a.equation = "ij,jk->ik";
        dims = RandomDims(2, 2);
      } else {
        a.equation = "...ij,...jk->...ik";
        dims = RandomDims(2);
      }
      int64_t ndims = dims.size();
      int64_t inner_dim = RandomDim();
      a.lhs_dims = dims;
      a.rhs_dims = dims;
      a.lhs_dims[ndims - 1] = inner_dim;
      a.rhs_dims[ndims - 2] = inner_dim;
      break;
    }
    case dot: {
      a.equation = "i,i->";
      std::vector<int64_t> dims = RandomDims(1, 1);
      a.lhs_dims = dims;
      a.rhs_dims = dims;
      break;
    }
    case outer: {
      a.equation = "i,j->ij";
      a.lhs_dims = RandomDims(1, 1);
      a.rhs_dims = RandomDims(1, 1);
      break;
    }
  }

  a.type = Choose<DataType>(kAllXlaTypes);
  return a;
}

OpTest::GatherArguments OpTest::ChooseGatherArguments(bool axis_0) {
  GatherArguments a;

  a.axis_type = DT_INT32;
  a.indices_type = DT_INT32;
  a.params_type = Choose<DataType>(kAllXlaTypes);

  // Choose parameters such that
  // 0 <= batch_dims <= axis < params.rank <= kDefaultMaxRank
  a.batch_dims = 0;
  int64_t axis;
  if (axis_0) {
    axis = 0;
  } else {
    std::uniform_int_distribution<int64_t> axis_distribution(
        a.batch_dims, kDefaultMaxRank - 1);
    axis = axis_distribution(generator());
  }
  a.axis = test::AsScalar<int32_t>((int32_t)axis);
  a.params_shape = RandomDims(axis + 1, kDefaultMaxRank, 1, 16);
  std::vector<int64_t> indices_shape = RandomDims(0, 3, 0, 16);
  a.indices = RandomBoundedTensor<int32_t>(
      DT_INT32, 0, a.params_shape[axis] - 1, false, indices_shape);

  return a;
}

OpTest::PadArguments OpTest::ChoosePadArguments() {
  PadArguments a;

  a.input_type = Choose<DataType>(kAllXlaTypes);
  a.input_shape = RandomDims();
  int input_rank = a.input_shape.size();

  a.paddings_type = Choose<DataType>({DT_INT32, DT_INT64});
  std::vector<int64_t> paddings_vec;
  for (int i = 0; i < input_rank; ++i) {
    std::uniform_int_distribution<int> pad_distribution(0, a.input_shape[i]);
    int pad_size = pad_distribution(generator());
    std::uniform_int_distribution<int> lower_distribution(0, pad_size);
    int low_pad_size = lower_distribution(generator());
    paddings_vec.push_back(low_pad_size);
    paddings_vec.push_back(pad_size - low_pad_size);
    a.input_shape[i] -= pad_size;
  }
  CHECK(
      a.paddings.CopyFrom(AsIntTensor(a.paddings_type, paddings_vec),
                          TensorShape({static_cast<int64_t>(input_rank), 2})));

  a.constant_values = RandomTensor(a.input_type, false, {});

  return a;
}

OpTest::ScatterArguments OpTest::ChooseScatterArguments() {
  ScatterArguments a;

  a.type = Choose<DataType>(kAllXlaTypes);
  a.indices_type = DT_INT32;
  a.shape = RandomDims(1, kDefaultMaxRank, 1);
  int rank = a.shape.size();
  std::uniform_int_distribution<int32_t> index_len_dist(1, rank);
  int index_len = index_len_dist(generator());
  std::vector<int64_t> indices_first = RandomDims(1, kDefaultMaxRank - 1, 1);
  std::vector<int64_t> indices_shape(indices_first);
  indices_shape.push_back(index_len);
  std::vector<int64_t> updates_shape(indices_first);
  for (int i = 0; i < rank - index_len; ++i) {
    updates_shape.push_back(a.shape[index_len + i]);
  }
  Tensor indices_lo(a.indices_type, TensorShape(indices_shape));
  test::FillFn<int32_t>(&indices_lo, [](int i) -> int32_t { return 0; });
  Tensor indices_hi(a.indices_type, TensorShape(indices_shape));
  test::FillFn<int32_t>(&indices_hi, [index_len, &a](int i) -> int32_t {
    int idx_dim = i % index_len;
    return a.shape[idx_dim] - 1;
  });
  a.indices = RandomBoundedTensor(a.indices_type, indices_lo, indices_hi);
  a.updates = RandomTensor(a.type, false, updates_shape);

  return a;
}

OpTest::SliceArguments OpTest::ChooseSliceArguments(bool neg_one_size) {
  SliceArguments a;

  a.type = Choose<DataType>(kAllXlaTypes);
  a.indices_type = DT_INT32;
  a.shape = RandomDims();
  int rank = a.shape.size();

  std::vector<int32_t> indices(rank);
  a.size.resize(rank);
  for (int i = 0; i < rank; ++i) {
    indices[i] =
        std::uniform_int_distribution<int32_t>(0, a.shape[i])(generator());
    int64_t low = neg_one_size ? -1 : 0;
    a.size[i] = std::uniform_int_distribution<int64_t>(
        low, a.shape[i] - indices[i])(generator());
  }
  a.indices = test::AsTensor<int32_t>(indices);

  return a;
}

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) {
    absl::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_t pad_dummy;
      s = GetWindowedOutputSize(d.input_dims[i], d.kernel_dims[i],
                                /*dilation_rate=*/1, d.stride_dims[i],
                                d.padding, &d.output_dims[i], &pad_dummy);
    } while (!s.ok());
  }
  return d;
}

OpTest::XlaDotArguments OpTest::ChooseXlaDotArguments() {
  std::vector<int64_t> batch_dims = RandomDims(0, 2);
  std::vector<int64_t> contracting_dims = RandomDims(0, 2);
  std::vector<int64_t> lhs_outer_dims = RandomDims(0, 2);
  std::vector<int64_t> rhs_outer_dims = RandomDims(0, 2);

  XlaDotArguments a;
  a.lhs_dims.insert(a.lhs_dims.end(), batch_dims.begin(), batch_dims.end());
  a.lhs_dims.insert(a.lhs_dims.end(), contracting_dims.begin(),
                    contracting_dims.end());
  a.lhs_dims.insert(a.lhs_dims.end(), lhs_outer_dims.begin(),
                    lhs_outer_dims.end());
  a.rhs_dims.insert(a.rhs_dims.end(), batch_dims.begin(), batch_dims.end());
  a.rhs_dims.insert(a.rhs_dims.end(), contracting_dims.begin(),
                    contracting_dims.end());
  a.rhs_dims.insert(a.rhs_dims.end(), rhs_outer_dims.begin(),
                    rhs_outer_dims.end());

  xla::DotDimensionNumbers dnums;
  for (auto i = 0; i < batch_dims.size(); ++i) {
    dnums.add_lhs_batch_dimensions(i);
    dnums.add_rhs_batch_dimensions(i);
  }
  for (auto i = 0; i < contracting_dims.size(); ++i) {
    dnums.add_lhs_contracting_dimensions(batch_dims.size() + i);
    dnums.add_rhs_contracting_dimensions(batch_dims.size() + i);
  }
  dnums.SerializeToString(&a.dnums_encoded);

  a.precision_config_encoded = "";

  a.dtype = Choose<DataType>(kAllXlaTypes);
  return a;
}

std::vector<int64_t> OpTest::ImageDims(
    TensorFormat format, int batch, int feature,
    const std::vector<int64_t>& spatial_dims) {
  std::vector<int64_t> 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;
    default:
      LOG(FATAL) << "Tensor format " << ToString(format) << " not supported.";
  }
  return dims;
}

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

// Functions for comparing tensors.

template <typename T>
double Abs(T x) {
  return std::fabs(x);
}

template <>
double Abs<complex64>(complex64 x) {
  return std::abs(x);
}

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 Abs(x - y) < atol + rtol * Abs(x);
}

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

template <typename T>
std::string Str(T x) {
  return absl::StrCat(x);
}
template <>
std::string Str<complex64>(complex64 x) {
  return absl::StrCat("(", x.real(), ", ", x.imag(), ")");
}

template <typename T>
absl::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(
          absl::StrCat(i, "-th tensor element isn't close: ", Str(Tx(i)),
                       " vs. ", Str(Ty(i)), ". x = ", x.DebugString(),
                       "y = ", y.DebugString(), "atol = ", atol,
                       " rtol = ", rtol, " tol = ", atol + rtol * Abs(Tx(i))));
    }
  }
  return absl::OkStatus();
}

template <typename T>
absl::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(absl::StrCat(
          i, "-th tensor element isn't equal: ", Str(Tx(i)), " vs. ",
          Str(Ty(i)), ". x = ", x.DebugString(), "y = ", y.DebugString()));
    }
  }
  return absl::OkStatus();
}

absl::Status TensorsAreEqualImplBfloat16(const Tensor& x, const Tensor& y) {
  auto Tx = x.flat<bfloat16>();
  auto Ty = y.flat<bfloat16>();
  for (int i = 0; i < Tx.size(); ++i) {
    if (Tx(i) != Ty(i)) {
      return absl::InvalidArgumentError(absl::StrCat(
          i, "-th tensor element isn't equal: ", static_cast<float>(Tx(i)),
          " vs. ", static_cast<float>(Ty(i)), ". x = ", x.DebugString(),
          "y = ", y.DebugString()));
    }
  }
  return absl::OkStatus();
}

// 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.
absl::Status TensorsAreClose(const Tensor& a, const Tensor& b, double atol,
                             double rtol) {
  if (a.dtype() != b.dtype()) {
    return absl::InvalidArgumentError(absl::StrCat(
        "Tensors have different types: ", DataTypeString(a.dtype()), " and ",
        DataTypeString(b.dtype())));
  }
  if (!a.IsSameSize(b)) {
    return absl::InvalidArgumentError(
        absl::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_COMPLEX64:
      return TensorsAreCloseImpl<complex64>(a, b, atol, rtol);
    case DT_INT32:
      return TensorsAreEqualImpl<int32_t>(a, b);
    case DT_INT64:
      return TensorsAreEqualImpl<int64_t>(a, b);
    case DT_BOOL:
      return TensorsAreEqualImpl<bool>(a, b);
    case DT_BFLOAT16:
      return TensorsAreEqualImplBfloat16(a, b);
    default:
      LOG(FATAL) << "Unexpected type : " << DataTypeString(a.dtype());
  }
}

OpTest::TestResult OpTest::ExpectTfAndXlaOutputsAreClose(
    const OpTestBuilder& builder, double atol, double rtol) {
  const std::vector<OpTestBuilder::InputDescription>& inputs = builder.inputs();
  std::vector<Tensor> input_tensors;
  input_tensors.reserve(inputs.size());
  for (const OpTestBuilder::InputDescription& input : inputs) {
    if (input.type == DT_INVALID) {
      input_tensors.push_back(input.tensor);
    } else {
      std::vector<int64_t> dims;
      if (input.has_dims) {
        dims = input.dims;
      } else {
        dims = RandomDims();
      }
      if (!TensorSizeIsOk(dims)) {
        VLOG(1) << "Input: " << input.type << " "
                << TensorShape(input.dims).DebugString();
        VLOG(1) << "Ignoring oversize dims.";
        return kInvalid;
      }
      input_tensors.push_back(
          RandomTensor(input.type, input.needs_unique_values, dims));
    }
    VLOG(1) << "Input: " << input_tensors.back().DebugString();
  }

  std::string reference_device =
      LocalDeviceToFullDeviceName(*tf_xla_reference_device_ptr);
  std::string test_device =
      LocalDeviceToFullDeviceName(*tf_xla_test_device_ptr);

  DeviceNameUtils::ParsedName parsed_name;
  if (!DeviceNameUtils::ParseLocalName(*tf_xla_test_device_ptr, &parsed_name)) {
    LOG(ERROR) << "Could not parse device name: " << *tf_xla_test_device_ptr;
    return kFatalError;
  }
  DeviceType test_device_type(parsed_name.type);
  ++num_tests_;

  GraphDef graph;
  std::vector<std::string> expected_inputs, test_inputs;
  std::vector<std::string> expected_fetches, test_fetches;
  absl::Status status = builder.BuildGraph(
      absl::StrCat("test", num_tests_, "_expected"), reference_device,
      /*use_jit=*/false, &graph, /*test_node_def=*/nullptr, &expected_inputs,
      &expected_fetches);
  if (!status.ok()) {
    LOG(ERROR) << "Expected graph construction failed: " << status;
    return kFatalError;
  }

  NodeDef* node_def;
  status = builder.BuildGraph(absl::StrCat("test", num_tests_, "_test"),
                              test_device, tf_xla_test_use_jit, &graph,
                              &node_def, &test_inputs, &test_fetches);
  if (!status.ok()) {
    LOG(ERROR) << "Test graph construction failed: " << status;
    return kFatalError;
  }

  // Check that there's a kernel corresponding to 'node_def' on the device under
  // test.
  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 kInvalid;
  }

  // Create a session with the corresponding graph.
  SessionOptions session_options;
  session_.reset(NewSession(session_options));
  status = session_->Create(graph);
  if (!status.ok()) {
    LOG(ERROR) << "Session::Create() failed: " << status;
    return kFatalError;
  }

  std::vector<std::pair<std::string, Tensor>> expected_feeds(
      expected_inputs.size());
  std::vector<std::pair<std::string, Tensor>> test_feeds(test_inputs.size());
  CHECK_EQ(input_tensors.size(), expected_inputs.size());
  CHECK_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";
  absl::Status s =
      session_->Run(expected_feeds, expected_fetches, {}, &expected_outputs);
  if (!s.ok()) {
    VLOG(1) << "Expected graph failed with status: " << s << ". Ignoring test";
    return kInvalid;
  }
  for (const Tensor& expected : expected_outputs) {
    VLOG(1) << "Expected: " << expected.DebugString();
  }

  VLOG(1) << "Running test graph";
  status = session_->Run(test_feeds, test_fetches, {}, &test_outputs);
  if (!status.ok()) {
    LOG(ERROR) << "Test graph failed: " << status;
    return kFatalError;
  }

  CHECK_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);

  return kOk;
}

TEST_F(OpTest, _EagerConst) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("_EagerConst").RandomInput(type).Attr("T", type));
  });
}

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

TEST_F(OpTest, Acos) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Acos")
            .Input(RandomBoundedTensor<float>(DT_FLOAT, -1, 1, false))
            .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Acosh) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Acosh").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

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

TEST_F(OpTest, AddN) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    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.RandomInput(type, shape);
    }
    return ExpectTfAndXlaOutputsAreClose(builder);
  });
}

TEST_F(OpTest, AddV2) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    auto dims = BroadcastableDims();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("AddV2")
                                             .RandomInput(type, dims.first)
                                             .RandomInput(type, dims.second)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, All) {
  Repeatedly([this]() {
    std::vector<int64_t> data_dims = RandomDims();
    Tensor indices = RandomReductionIndices(data_dims.size());
    bool keep_dims = Choose<bool>({false, true});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("All")
                                             .RandomInput(DT_BOOL, data_dims)
                                             .Input(indices)
                                             .Attr("keep_dims", keep_dims));
  });
}

TEST_F(OpTest, Angle) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Angle")
                                             .RandomInput(DT_COMPLEX64)
                                             .Attr("T", DT_COMPLEX64));
  });
}

TEST_F(OpTest, Any) {
  Repeatedly([this]() {
    std::vector<int64_t> data_dims = RandomDims();
    Tensor indices = RandomReductionIndices(data_dims.size());
    bool keep_dims = Choose<bool>({false, true});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Any")
                                             .RandomInput(DT_BOOL, data_dims)
                                             .Input(indices)
                                             .Attr("keep_dims", keep_dims));
  });
}

TEST_F(OpTest, ApproximateEqual) {
  Repeatedly([this]() {
    auto dims = BroadcastableDims();
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ApproximateEqual")
                                             .RandomInput(type, dims.first)
                                             .RandomInput(type, dims.second)
                                             .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, ArgMax) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_BOOL, DT_FLOAT});
    std::vector<int64_t> dims = RandomDims(1, 5, 1);
    int num_dims = dims.size();
    int reduce_dim = std::uniform_int_distribution<int32_t>(
        -num_dims, num_dims)(generator());
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("ArgMax")
            .RandomInput(type, dims)
            .Input(test::AsScalar<int32_t>(reduce_dim))
            .Attr("T", type)
            .Attr("Tidx", DT_INT32)
            .Attr("output_type", DT_INT32));
  });
}

TEST_F(OpTest, ArgMin) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_BOOL, DT_FLOAT});
    std::vector<int64_t> dims = RandomDims(1, 5, 1);
    int num_dims = dims.size();
    int reduce_dim = std::uniform_int_distribution<int32_t>(
        -num_dims, num_dims)(generator());
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("ArgMin")
            .RandomInput(type, dims)
            .Input(test::AsScalar<int32_t>(reduce_dim))
            .Attr("T", type)
            .Attr("Tidx", DT_INT32)
            .Attr("output_type", DT_INT32));
  });
}

TEST_F(OpTest, Asin) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Asin")
            .Input(RandomBoundedTensor<float>(DT_FLOAT, -1, 1, false))
            .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Asinh) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Asinh").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Atanh) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Atanh").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Atan) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Atan").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

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

TEST_F(OpTest, AvgPool) {
  Repeatedly([this]() {
    std::uniform_int_distribution<int> random_int(1, 5);
    std::vector<int64_t> dims = RandomDims(4, 4, 1);
    int kernel_rows =
        std::uniform_int_distribution<int>(1, dims[1])(generator());
    int kernel_cols =
        std::uniform_int_distribution<int>(1, dims[2])(generator());
    int stride_rows = random_int(generator()),
        stride_cols = random_int(generator());
    std::string padding = Choose<std::string>({"SAME", "VALID"});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("AvgPool")
            .RandomInput(DT_FLOAT, dims)
            .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) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    std::uniform_int_distribution<int> random_int(1, 5);
    std::vector<int64_t> dims = RandomDims(5, 5, 1);

    std::vector<int64_t> input_dims, kernel_dims, stride_dims;
    for (int i = 0; i < 3; ++i) {
      kernel_dims.push_back(
          std::uniform_int_distribution<int>(1, dims[i])(generator()));
      input_dims.push_back(dims[i]);
      stride_dims.push_back(random_int(generator()));
    }
    int64_t batch = dims[3];
    int64_t feature = dims[4];

    std::string padding = Choose<std::string>({"SAME", "VALID"});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("AvgPool3D")
            .RandomInput(DT_FLOAT,
                         ImageDims(FORMAT_NHWC, batch, feature, 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", "NDHWC"));
  });
  // TODO(phawkins): test NCHW format (not supported by CPU)
}

TEST_F(OpTest, AvgPoolGrad) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    int batch = RandomDim(1), features = RandomDim(1);
    WindowedSpatialDims d = ChooseWindowedSpatialDims(2);
    std::vector<int32_t> input_dims =
        AsInt32s(ImageDims(FORMAT_NHWC, batch, features, d.input_dims));
    std::vector<int64_t> output_dims =
        ImageDims(FORMAT_NHWC, batch, features, d.output_dims);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("AvgPoolGrad")
            .Input(test::AsTensor<int32_t>(input_dims))
            .RandomInput(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) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    int batch = RandomDim(1), features = RandomDim(1);
    WindowedSpatialDims d = ChooseWindowedSpatialDims(3);
    std::vector<int32_t> input_dims =
        AsInt32s(ImageDims(FORMAT_NHWC, batch, features, d.input_dims));
    std::vector<int64_t> output_dims =
        ImageDims(FORMAT_NHWC, batch, features, d.output_dims);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("AvgPool3DGrad")
            .Input(test::AsTensor<int32_t>(input_dims))
            .RandomInput(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", "NDHWC"));
  });
}

TEST_F(OpTest, BatchMatMul) {
  // See note about failing Kokoro tests: b/214080339#comment22
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    const BatchMatMulArguments a = ChooseBatchMatMulArguments(false);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BatchMatMul")
                                             .RandomInput(a.dtype, a.lhs_dims)
                                             .RandomInput(a.dtype, a.rhs_dims)
                                             .Attr("T", a.dtype)
                                             .Attr("adj_x", a.adj_lhs)
                                             .Attr("adj_y", a.adj_rhs));
  });
}

TEST_F(OpTest, BatchMatMulV2) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  // :randomized_tests_seeded is flaky with --tf_xla_random_seed=200839030
  // See b/229622638.
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    const BatchMatMulArguments a = ChooseBatchMatMulArguments(true);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BatchMatMulV2")
                                             .RandomInput(a.dtype, a.lhs_dims)
                                             .RandomInput(a.dtype, a.rhs_dims)
                                             .Attr("T", a.dtype)
                                             .Attr("adj_x", a.adj_lhs)
                                             .Attr("adj_y", a.adj_rhs));
  });
}

TEST_F(OpTest, BatchToSpace) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    const int num_block_dims = 2;
    std::vector<int64_t> block_dims =
        RandomDims(num_block_dims, num_block_dims, 0, 5);
    int64_t block_size = RandomDim(2, 5);

    std::vector<int64_t> 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_t> 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})));

    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BatchToSpace")
                                             .RandomInput(type, input_dims)
                                             .Input(crops)
                                             .Attr("T", type)
                                             .Attr("block_size", block_size));
  });
}

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

    std::vector<int64_t> 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_t> 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})));

    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("BatchToSpaceND")
            .RandomInput(type, input_dims)
            .Input(test::AsTensor<int32_t>(
                std::vector<int32_t>(block_dims.begin(), block_dims.end())))
            .Input(crops)
            .Attr("T", type));
  });
}

TEST_F(OpTest, BiasAdd) {
  Repeatedly([this]() {
    auto x_dims = RandomDims(2, kDefaultMaxRank);
    auto y_dims = {x_dims[x_dims.size() - 1]};
    // TODO(phawkins): test both data formats.
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BiasAdd")
                                             .RandomInput(type, x_dims)
                                             .RandomInput(type, y_dims)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, BiasAddGrad) {
  Repeatedly([this]() {
    // TODO(phawkins): test both data formats.
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("BiasAddGrad").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, BiasAddV1) {
  Repeatedly([this]() {
    auto x_dims = RandomDims(2, kDefaultMaxRank);
    auto y_dims = {x_dims[x_dims.size() - 1]};
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BiasAddV1")
                                             .RandomInput(type, x_dims)
                                             .RandomInput(type, y_dims)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, Bitcast) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  Repeatedly([this]() {  // NOLINT: due to GTEST_SKIP
    auto src_type = Choose<DataType>(kAllNumberTypes);
    auto dst_type = Choose<DataType>(kAllNumberTypes);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Bitcast")
                                             .RandomInput(src_type)
                                             .Attr("T", src_type)
                                             .Attr("type", dst_type));
  });
}

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

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

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

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

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

TEST_F(OpTest, BroadcastTo) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    auto type_idx = Choose<DataType>({DT_INT32, DT_INT64});
    auto dims_to = RandomDims();
    auto dims_from = BroadcastableToDims(dims_to);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("BroadcastTo")
            .RandomInput(type, dims_from)
            .Input(AsIntTensor(type_idx, dims_to))
            .Attr("T", type)
            .Attr("Tidx", type_idx));
  });
}

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

TEST_F(OpTest, CastBF16) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  Repeatedly([this]() {
    DataType src_type, dst_type;
    src_type = Choose<DataType>({DT_FLOAT});
    dst_type = Choose<DataType>({DT_BFLOAT16});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Cast")
                                             .RandomInput(src_type)
                                             .Attr("SrcT", src_type)
                                             .Attr("DstT", dst_type)
                                             .Attr("Truncate", true));
  });
}

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

TEST_F(OpTest, ClipByValue) {
  // TODO(b/211012085): Change input_dims to BroadcastableDimsN(3). The
  //                    compiled ClipByValue fails in this case.
  //                    --tf_xla_random_seed=200839030
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_INT32, DT_INT64, DT_FLOAT});
    // ClipByValue requires that broadcasting min and max tensors do not cause
    // the returned shape to be larger than the input shape.
    auto input_dims = RandomDims();
    // clip_value_min must be <= clip_value_max for correct results. Different
    // implementations handle the max < min case differently, so ensure that
    // min <= max.
    auto min_max_dims = BroadcastableToDims(input_dims);
    auto min_max = RandomLteTensors(type, min_max_dims);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ClipByValue")
                                             .RandomInput(type, input_dims)
                                             .Input(min_max.first)
                                             .Input(min_max.second)
                                             .Attr("T", type));
  });
}

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

TEST_F(OpTest, Concat) {
  Repeatedly([this]() {  // NOLINT: due to GTEST_SKIP
    ConcatArguments a = ChooseConcatArguments(false);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Concat")
                                             .Input(a.axis)
                                             .VariadicInput(a.values)
                                             .Attr("N", a.n)
                                             .Attr("T", a.type));
  });
}

TEST_F(OpTest, ConcatV2) {
  Repeatedly([this]() {
    ConcatArguments a = ChooseConcatArguments(true);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ConcatV2")
                                             .VariadicInput(a.values)
                                             .Input(a.axis)
                                             .Attr("N", a.n)
                                             .Attr("T", a.type)
                                             .Attr("Tidx", a.type_idx));
  });
}

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

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

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

TEST_F(OpTest, Conj) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Conj")
                                             .RandomInput(DT_COMPLEX64)
                                             .Attr("T", DT_COMPLEX64));
  });
}

TEST_F(OpTest, Const) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Const")
                                             .Attr("value", RandomTensor(type))
                                             .Attr("dtype", type));
  });
}

TEST_F(OpTest, FFT) {
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims(1, kDefaultMaxRank);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("FFT").RandomInput(DT_COMPLEX64, dims));
  });
}

TEST_F(OpTest, FFT2D) {
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims(2, kDefaultMaxRank);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("FFT2D").RandomInput(DT_COMPLEX64, dims));
  });
}

TEST_F(OpTest, FFT3D) {
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims(3, kDefaultMaxRank);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("FFT3D").RandomInput(DT_COMPLEX64, dims));
  });
}

TEST_F(OpTest, IFFT) {
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims(1, kDefaultMaxRank);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("IFFT").RandomInput(DT_COMPLEX64, dims));
  });
}

TEST_F(OpTest, IFFT2D) {
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims(2, kDefaultMaxRank);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("IFFT2D").RandomInput(DT_COMPLEX64, dims));
  });
}

TEST_F(OpTest, IFFT3D) {
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims(3, kDefaultMaxRank);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("IFFT3D").RandomInput(DT_COMPLEX64, dims));
  });
}

TEST_F(OpTest, RFFT) {
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims(1, kDefaultMaxRank, 3);
    Tensor fft_shape =
        test::AsTensor<int32_t>(AsInt32s({dims[dims.size() - 1]}));
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("RFFT").RandomInput(DT_FLOAT, dims).Input(fft_shape));
  });
}

TEST_F(OpTest, RFFT2D) {
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims(2, kDefaultMaxRank, 3);
    Tensor fft_shape = test::AsTensor<int32_t>(
        AsInt32s({dims[dims.size() - 2], dims[dims.size() - 1]}));
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("RFFT2D").RandomInput(DT_FLOAT, dims).Input(fft_shape));
  });
}

TEST_F(OpTest, RFFT3D) {
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims(3, kDefaultMaxRank, 3);
    Tensor fft_shape = test::AsTensor<int32_t>(AsInt32s(
        {dims[dims.size() - 3], dims[dims.size() - 2], dims[dims.size() - 1]}));
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("RFFT3D").RandomInput(DT_FLOAT, dims).Input(fft_shape));
  });
}

TEST_F(OpTest, IRFFT) {
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims(1, kDefaultMaxRank, 3);
    int64_t orig_size = dims[dims.size() - 1];
    dims[dims.size() - 1] = dims[dims.size() - 1] / 2 + 1;
    Tensor fft_shape = test::AsTensor<int32_t>(AsInt32s({orig_size}));
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("IRFFT")
                                             .RandomInput(DT_COMPLEX64, dims)
                                             .Input(fft_shape));
  });
}

TEST_F(OpTest, IRFFT2D) {
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims(2, kDefaultMaxRank, 3);
    std::vector<int64_t> orig_size = {dims[dims.size() - 2],
                                      dims[dims.size() - 1]};
    dims[dims.size() - 1] = dims[dims.size() - 1] / 2 + 1;
    Tensor fft_shape = test::AsTensor<int32_t>(AsInt32s({orig_size}));
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("IRFFT2D")
                                             .RandomInput(DT_COMPLEX64, dims)
                                             .Input(fft_shape));
  });
}

TEST_F(OpTest, IRFFT3D) {
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims(3, kDefaultMaxRank, 3);
    std::vector<int64_t> orig_size = {
        dims[dims.size() - 3], dims[dims.size() - 2], dims[dims.size() - 1]};
    dims[dims.size() - 1] = dims[dims.size() - 1] / 2 + 1;
    Tensor fft_shape = test::AsTensor<int32_t>(AsInt32s({orig_size}));
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("IRFFT3D")
                                             .RandomInput(DT_COMPLEX64, dims)
                                             .Input(fft_shape));
  });
}

TEST_F(OpTest, Conv2D) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  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_t batch = RandomDim();

    std::vector<int64_t> data_dims =
        ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims);

    std::vector<int64_t> kernel_dims = {d.kernel_dims[0], d.kernel_dims[1],
                                        features_in, features_out};
    DataType type = DT_FLOAT;
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Conv2D")
            .RandomInput(type, data_dims)
            .RandomInput(type, kernel_dims)
            .Attr("T", type)
            .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_t batch = RandomDim();
    std::vector<int64_t> activations =
        ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims);
    std::vector<int64_t> backprop =
        ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims);
    Tensor kernel_shape = test::AsTensor<int32_t>(AsInt32s(
        {d.kernel_dims[0], d.kernel_dims[1], features_in, features_out}));
    DataType type = DT_FLOAT;
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Conv2DBackpropFilter")
            .RandomInput(type, activations)
            .Input(kernel_shape)
            .RandomInput(type, backprop)
            .Attr("T", type)
            .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_t batch = RandomDim();
    Tensor in_shape = test::AsTensor<int32_t>(
        AsInt32s(ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims)));
    std::vector<int64_t> backprop =
        ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims);
    std::vector<int64_t> kernel = {d.kernel_dims[0], d.kernel_dims[1],
                                   features_in, features_out};
    DataType type = DT_FLOAT;
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Conv2DBackpropInput")
            .Input(in_shape)
            .RandomInput(type, kernel)
            .RandomInput(type, backprop)
            .Attr("T", type)
            .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) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  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());
    std::vector<int64_t> data = {RandomDim(), d.input_dims[0], d.input_dims[1],
                                 d.input_dims[2], features_in};

    std::vector<int64_t> kernel = {d.kernel_dims[0], d.kernel_dims[1],
                                   d.kernel_dims[2], features_in, features_out};
    DataType type = DT_FLOAT;
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Conv3D")
            .RandomInput(type, data)
            .RandomInput(type, kernel)
            .Attr("T", type)
            .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims))
            .Attr("padding", d.padding == SAME ? "SAME" : "VALID"));
  });
}

TEST_F(OpTest, Conv3DBackpropFilter) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  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_t batch = RandomDim(1);
    std::vector<int64_t> activations =
        ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims);
    std::vector<int64_t> backprop =
        ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims);
    Tensor kernel_shape = test::AsTensor<int32_t>(
        AsInt32s({d.kernel_dims[0], d.kernel_dims[1], d.kernel_dims[2],
                  features_in, features_out}));
    DataType type = DT_FLOAT;
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Conv3DBackpropFilterV2")
            .RandomInput(type, activations)
            .Input(kernel_shape)
            .RandomInput(type, backprop)
            .Attr("T", type)
            .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims))
            .Attr("padding", d.padding == SAME ? "SAME" : "VALID"));
  });
}

TEST_F(OpTest, Conv3DBackpropInput) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  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_t batch = RandomDim(1);
    Tensor in_shape = test::AsTensor<int32_t>(
        AsInt32s(ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims)));
    std::vector<int64_t> backprop =
        ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims);
    std::vector<int64_t> kernel = {d.kernel_dims[0], d.kernel_dims[1],
                                   d.kernel_dims[2], features_in, features_out};
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Conv3DBackpropInputV2")
            .Input(in_shape)
            .RandomInput(type, kernel)
            .RandomInput(type, backprop)
            .Attr("T", type)
            .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims))
            .Attr("padding", d.padding == SAME ? "SAME" : "VALID"));
  });
}

TEST_F(OpTest, ComplexAbs) {
  Repeatedly([this]() {
    auto type = DT_COMPLEX64;
    auto type_out = DT_FLOAT;
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ComplexAbs")
                                             .RandomInput(type)
                                             .Attr("T", type)
                                             .Attr("Tout", type_out));
  });
}

TEST_F(OpTest, Cos) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Cos").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, Cosh) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Cosh").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, DepthToSpace) {
  Repeatedly([this]() {
    int64_t block = RandomDim(2, 5);
    std::vector<int64_t> input_dims = RandomDims(4, 4);
    input_dims[1] = (input_dims[1] + (block - 1)) / block;
    input_dims[2] = (input_dims[2] + (block - 1)) / block;
    input_dims[3] *= block * block;
    auto type = Choose<DataType>(kAllXlaTypes);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("DepthToSpace")
                                             .RandomInput(type, input_dims)
                                             .Attr("T", type)
                                             .Attr("block_size", block));
  });
}

TEST_F(OpTest, DepthwiseConv2DNative) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    WindowedSpatialDims d = ChooseWindowedSpatialDims(2);
    std::uniform_int_distribution<int> random_int(1, 5);
    int features_in = random_int(generator());
    int depth_multiplier = random_int(generator());
    std::vector<int64_t> input_dims = {RandomDim(), d.input_dims[0],
                                       d.input_dims[1], features_in};

    std::vector<int64_t> kernel_dims = {d.kernel_dims[0], d.kernel_dims[1],
                                        features_in, depth_multiplier};
    std::vector<int64_t> strides = ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims);
    strides[2] = strides[1];  // Current impl only supports equal strides
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("DepthwiseConv2dNative")
            .RandomInput(DT_FLOAT, input_dims)
            .RandomInput(DT_FLOAT, kernel_dims)
            .Attr("T", DT_FLOAT)
            .Attr("strides", strides)
            .Attr("padding", d.padding == SAME ? "SAME" : "VALID"));
  });
}

TEST_F(OpTest, DepthwiseConv2DNativeBackpropFilter) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    WindowedSpatialDims d = ChooseWindowedSpatialDims(2);
    std::uniform_int_distribution<int> random_int(1, 5);
    int features_in = random_int(generator());
    int depth_multiplier = random_int(generator());
    int32_t batch = RandomDim();
    std::vector<int64_t> activations =
        ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims);
    std::vector<int64_t> backprop = ImageDims(
        FORMAT_NHWC, batch, features_in * depth_multiplier, d.output_dims);
    Tensor kernel_shape = test::AsTensor<int32_t>(AsInt32s(
        {d.kernel_dims[0], d.kernel_dims[1], features_in, depth_multiplier}));
    std::vector<int64_t> strides = ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims);
    strides[2] = strides[1];  // Current impl only supports equal strides
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("DepthwiseConv2dNativeBackpropFilter")
            .RandomInput(DT_FLOAT, activations)
            .Input(kernel_shape)
            .RandomInput(DT_FLOAT, backprop)
            .Attr("T", DT_FLOAT)
            .Attr("strides", strides)
            .Attr("padding", d.padding == SAME ? "SAME" : "VALID")
            .Attr("data_format", "NHWC"));
  });
}

TEST_F(OpTest, DepthwiseConv2DBackpropInput) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    WindowedSpatialDims d = ChooseWindowedSpatialDims(2);
    std::uniform_int_distribution<int> random_int(1, 5);
    int features_in = random_int(generator());
    int depth_multiplier = random_int(generator());
    int32_t batch = RandomDim();
    Tensor in_shape = test::AsTensor<int32_t>(
        AsInt32s(ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims)));
    std::vector<int64_t> backprop = ImageDims(
        FORMAT_NHWC, batch, features_in * depth_multiplier, d.output_dims);
    std::vector<int64_t> kernel = {d.kernel_dims[0], d.kernel_dims[1],
                                   features_in, depth_multiplier};
    std::vector<int64_t> strides = ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims);
    strides[2] = strides[1];  // Current impl only supports equal strides
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("DepthwiseConv2dNativeBackpropInput")
            .Input(in_shape)
            .RandomInput(DT_FLOAT, kernel)
            .RandomInput(DT_FLOAT, backprop)
            .Attr("T", DT_FLOAT)
            .Attr("strides", strides)
            .Attr("padding", d.padding == SAME ? "SAME" : "VALID")
            .Attr("data_format", "NHWC"));
  });
}

TEST_F(OpTest, Diag) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    std::vector<int64_t> dims;
    // Diag causes a quadratic blowup in output size.
    int64_t size;
    do {
      dims = RandomDims(1);
      size = TensorShape(dims).num_elements();
    } while (size * size > tf_xla_max_tensor_size);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Diag").RandomInput(type, dims).Attr("T", type));
  });
}

TEST_F(OpTest, DiagPart) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    auto dims = RandomDims(1, 3);
    // Duplicate the random dims.
    std::vector<int64_t> 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());
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("DiagPart")
                                             .RandomInput(type, doubled_dims)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, Digamma) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Digamma").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

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

TEST_F(OpTest, DivNoNan) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    auto dims = BroadcastableDims();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("DivNoNan")
                                             .RandomInput(type, dims.first)
                                             .RandomInput(type, dims.second)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, DynamicStitch) {
  Repeatedly([this]() {
    auto 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_t>> 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_t> 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_t> 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_t>(
          absl::Span<const int32_t>(indices).subspan(pos, shape.num_elements()),
          shape);
      builder.Input(t);
      pos += t.NumElements();
    }

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

TEST_F(OpTest, Einsum) {
  Repeatedly([this]() {
    const EinsumArguments a = ChooseEinsumArguments();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Einsum")
                                             .RandomInput(a.type, a.lhs_dims)
                                             .RandomInput(a.type, a.rhs_dims)
                                             .Attr("equation", a.equation)
                                             .Attr("T", a.type)
                                             .Attr("N", 2),
                                         2e-1, 2e-1);
  });
}

TEST_F(OpTest, Empty) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({kAllXlaTypes});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Empty")
            .Input(AsIntTensor(DT_INT32, RandomDims()))
            .Attr("init", true)
            .Attr("dtype", type));
  });
}

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

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

TEST_F(OpTest, ScatterNd) {
  Repeatedly([this]() {
    auto a = ChooseScatterArguments();
    auto shape = test::AsTensor<int32_t>(
        std::vector<int32_t>(a.shape.begin(), a.shape.end()));
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ScatterNd")
                                             .Input(a.indices)
                                             .Input(a.updates)
                                             .Input(shape)
                                             .Attr("T", a.type)
                                             .Attr("Tindices", a.indices_type));
  });
}

TEST_F(OpTest, Selu) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Selu").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, SeluGrad) {
  Repeatedly([this]() {
    auto dims = RandomDims();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SeluGrad")
                                             .RandomInput(DT_FLOAT, dims)
                                             .RandomInput(DT_FLOAT, dims)
                                             .Attr("T", DT_FLOAT));
  });
}

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

TEST_F(OpTest, Erf) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Erf").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Erfc) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Erfc").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Exp) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Exp").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, Expm1) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Expm1").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, ExpandDims) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    std::vector<int64_t> in_dims = RandomDims();
    Tensor dim(DT_INT32, TensorShape());
    std::uniform_int_distribution<int32_t> d(-1 - in_dims.size(),
                                             in_dims.size());
    dim.scalar<int32_t>()() = d(generator());
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ExpandDims")
                                             .RandomInput(type, in_dims)
                                             .Input(dim)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, Fill) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    std::vector<int64_t> dims = RandomDims();
    std::vector<int32_t> shape(dims.begin(), dims.end());
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Fill")
            .Input(test::AsTensor<int32_t>(shape))
            .RandomInput(type, {})
            .Attr("T", type));
  });
}

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

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

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

TEST_F(OpTest, Gather) {
  Repeatedly([this]() {
    GatherArguments a = ChooseGatherArguments(true);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Gather")
            .RandomInput(a.params_type, a.params_shape)
            .Input(a.indices)
            .Attr("Tparams", a.params_type)
            .Attr("Tindices", a.indices_type));
  });
}

TEST_F(OpTest, GatherV2) {
  Repeatedly([this]() {
    GatherArguments a = ChooseGatherArguments(false);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("GatherV2")
            .RandomInput(a.params_type, a.params_shape)
            .Input(a.indices)
            .Input(a.axis)
            .Attr("batch_dims", a.batch_dims)
            .Attr("Taxis", a.axis_type)
            .Attr("Tindices", a.indices_type)
            .Attr("Tparams", a.params_type));
  });
}

TEST_F(OpTest, GatherNd) {
  // :randomized_tests_mlir fails with --tf_xla_random_seed=459353625
  // --test_arg=--tf_xla_test_repetitions=100
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  // See b/214080339#comment27 as this test causes Kokoro to crash.
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {  // NOLINT: due to GTEST_SKIP
    auto params_type = Choose<DataType>(kAllXlaTypes);
    // GatherNd seems undefined on the case where params has rank 0.
    std::vector<int64_t> params_shape = RandomDims(1);
    auto indices_type = DT_INT32;
    std::vector<int64_t> output_outer_shape = RandomDims(0, 4, 0, 32);
    int64_t index_len = RandomDim(0, params_shape.size() + 1);
    std::vector<int64_t> output_shape(output_outer_shape);
    output_shape.push_back(index_len);
    Tensor lo(indices_type, TensorShape(output_shape));
    test::FillFn<int32_t>(&lo, [](int i) -> int32_t { return 0; });
    Tensor hi(indices_type, TensorShape(output_shape));
    test::FillFn<int32_t>(&hi, [index_len, &params_shape](int i) -> int32_t {
      int idx_dim = i % index_len;
      return params_shape[idx_dim] - 1;
    });
    Tensor indices = RandomBoundedTensor(indices_type, lo, hi);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("GatherNd")
            .RandomInput(params_type, params_shape)
            .Input(indices)
            .Attr("Tindices", indices_type)
            .Attr("Tparams", params_type));
  });
}

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

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

TEST_F(OpTest, Identity) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Identity").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, Imag) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Imag")
                                             .RandomInput(DT_COMPLEX64)
                                             .Attr("T", DT_COMPLEX64));
  });
}

TEST_F(OpTest, InplaceUpdate) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    std::vector<int64_t> common_dims =
        RandomDims(0, kDefaultMaxRank - 1, 0, kDefaultMaxDimensionSize);
    // TODO(b/211012712): Once needs_unique_values case is linear instead of
    // quadratic time, use default Dim max instead of 8.
    std::vector<int64_t> v_dims{RandomDim(1, 8)};
    v_dims.insert(v_dims.end(), common_dims.begin(), common_dims.end());
    std::vector<int64_t> x_dims{RandomDim(v_dims[0])};
    x_dims.insert(x_dims.end(), common_dims.begin(), common_dims.end());
    std::vector<int64_t> i_shape{v_dims[0]};
    Tensor i =
        RandomBoundedTensor<int32_t>(DT_INT32, 0, x_dims[0] - 1, true, i_shape);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("InplaceUpdate")
                                             .RandomInput(type, x_dims)
                                             .Input(i)
                                             .RandomInput(type, v_dims)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, Inv) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Inv").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, Invert) {
  Repeatedly([this]() {
    DataType type = DT_INT32;
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Invert").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, InvertPermutation) {
  Repeatedly([this]() {
    // TODO(b/211012712): Once needs_unique_values case is linear instead of
    // quadratic time, use default Dim max instead of 8.
    int64_t len = RandomDim(0, 8);
    Tensor x = RandomBoundedTensor<int32_t>(DT_INT32, 0, len - 1, true, {len});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("InvertPermutation").Input(x).Attr("T", DT_INT32));
  });
}

TEST_F(OpTest, IsFinite) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("IsFinite").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, IsInf) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("IsInf").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, IsNan) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("IsNan").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, L2Loss) {
  Repeatedly([this]() {
    DataType type = DT_FLOAT;
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("L2Loss").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, LeakyRelu) {
  Repeatedly([this]() {
    std::uniform_real_distribution<float> alpha(-2.0f, 2.0f);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("LeakyRelu")
            .RandomInput(DT_FLOAT)
            .Attr("T", DT_FLOAT)
            .Attr("alpha", alpha(generator())));
  });
}

TEST_F(OpTest, LeakyReluGrad) {
  Repeatedly([this]() {
    auto dims = RandomDims(1);
    std::uniform_real_distribution<float> alpha(-2.0f, 2.0f);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("LeakyReluGrad")
            .RandomInput(DT_FLOAT, dims)
            .RandomInput(DT_FLOAT, dims)
            .Attr("T", DT_FLOAT)
            .Attr("alpha", alpha(generator())));
  });
}

TEST_F(OpTest, LeftShift) {
  Repeatedly([this]() {
    bool is64 = RandomBool();
    auto dims = RandomDims();
    auto type = is64 ? DT_INT64 : DT_INT32;
    int max_shift = is64 ? 63 : 31;
    auto y = RandomBoundedTensor(type, 0, max_shift, false, dims);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("LeftShift")
                                             .RandomInput(type, dims)
                                             .Input(y)
                                             .Attr("T", type));
  });
}

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

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

TEST_F(OpTest, Lgamma) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Lgamma").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

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

TEST_F(OpTest, Log) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Log").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, Log1p) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Log1p").RandomInput(type).Attr("T", DT_FLOAT));
  });
}

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

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

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

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

TEST_F(OpTest, LRN) {
  Repeatedly([this]() {
    // TODO(b/31362467): Crashes with 0 dims on GPU. Re-enable when fixed.
    std::vector<int64_t> data_dims = RandomDims(4, 4, 1, 8);
    // CuDNN requires depth_radius > 0.
    std::uniform_int_distribution<int> radius(1, data_dims[3]);
    std::uniform_real_distribution<float> coeff(0.01, 2.0);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("LRN")
            .RandomInput(DT_FLOAT, data_dims)
            .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_t> dims = RandomDims(4, 4, 1, 8);
    // CuDNN requires depth_radius > 0.
    std::uniform_int_distribution<int> radius(1, dims[3]);
    std::uniform_real_distribution<float> coeff(0.0, 2.0);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("LRNGrad")
            .RandomInput(DT_FLOAT, dims)
            .RandomInput(DT_FLOAT, dims)
            .RandomInput(DT_FLOAT, dims)
            .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_t x = RandomDim();
    int64_t y = RandomDim();
    int64_t z = RandomDim();

    std::vector<int64_t> a_dims = {x, y};
    std::vector<int64_t> b_dims = {y, z};

    std::bernoulli_distribution random_bool;
    bool transpose_a = random_bool(generator());
    bool transpose_b = random_bool(generator());
    if (transpose_a) {
      std::swap(a_dims[0], a_dims[1]);
    }
    if (transpose_b) {
      std::swap(b_dims[0], b_dims[1]);
    }

    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatMul")
                                             .RandomInput(type, a_dims)
                                             .RandomInput(type, b_dims)
                                             .Attr("T", type)
                                             .Attr("transpose_a", transpose_a)
                                             .Attr("transpose_b", transpose_b));
  });
}

TEST_F(OpTest, MatrixBandPart) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    auto index_type = Choose<DataType>({DT_INT32, DT_INT64});
    auto num_lower =
        RandomBoundedTensor<int32_t>(index_type, -2 * kDefaultMaxDimensionSize,
                                     2 * kDefaultMaxDimensionSize, false, {});
    auto num_upper =
        RandomBoundedTensor<int32_t>(index_type, -2 * kDefaultMaxDimensionSize,
                                     2 * kDefaultMaxDimensionSize, false, {});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatrixBandPart")
                                             .RandomInput(type)
                                             .Input(num_lower)
                                             .Input(num_upper)
                                             .Attr("T", type)
                                             .Attr("Tindex", index_type));
  });
}

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

TEST_F(OpTest, MatrixDiagPart) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatrixDiagPart")
                                             .RandomInput(type, RandomDims(2))
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, MatrixDiagPartV3) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {  // NOLINT: due to GTEST_SKIP
    auto type = Choose<DataType>(kAllXlaTypes);
    auto align = Choose<std::string>(
        {"LEFT_RIGHT", "RIGHT_LEFT", "LEFT_LEFT", "RIGHT_RIGHT"});
    auto k0 = std::uniform_int_distribution<int32_t>(
        -2 * kDefaultMaxDimensionSize,
        2 * kDefaultMaxDimensionSize)(generator());
    auto k1 = std::uniform_int_distribution<int32_t>(
        k0, 2 * kDefaultMaxDimensionSize)(generator());
    auto k = test::AsTensor<int32_t>({k0, k1});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatrixDiagPartV3")
                                             .RandomInput(type)
                                             .Input(k)
                                             .RandomInput(type, {})
                                             .Attr("align", align)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, MatrixSetDiag) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    auto shape = RandomDims(2);
    int rank = shape.size();
    std::vector<int64_t> diagonal_shape(shape);
    diagonal_shape.pop_back();
    diagonal_shape.pop_back();
    diagonal_shape.push_back(std::min(shape[rank - 2], shape[rank - 1]));
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatrixSetDiag")
                                             .RandomInput(type, shape)
                                             .RandomInput(type, diagonal_shape)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, MatrixSetDiagV2) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    auto shape = RandomDims(2, kDefaultMaxRank, 1 /* non-zero dims */);
    int rank = shape.size();
    int64_t max_num_diags = shape[rank - 2] + shape[rank - 1] - 1;
    int64_t num_diags =
        std::uniform_int_distribution<int64_t>(2, max_num_diags)(generator());
    int32_t k0 = std::uniform_int_distribution<int32_t>(
        -shape[rank - 2] + 1, shape[rank - 1] - num_diags)(generator());
    int32_t k1 = k0 + num_diags - 1;
    Tensor k = test::AsTensor<int32_t>({k0, k1});
    int64_t max_diag_len = std::min(shape[rank - 2] + std::min(k1, 0),
                                    shape[rank - 1] + std::min(-k0, 0));
    std::vector<int64_t> diagonal_shape(shape);
    diagonal_shape.pop_back();
    diagonal_shape.pop_back();
    diagonal_shape.push_back(num_diags);
    diagonal_shape.push_back(max_diag_len);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatrixSetDiagV2")
                                             .RandomInput(type, shape)
                                             .RandomInput(type, diagonal_shape)
                                             .Input(k)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, Max) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT});
    std::vector<int64_t> data_dims = RandomDims();
    Tensor indices = RandomReductionIndices(data_dims.size());
    bool keep_dims = Choose<bool>({false, true});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Max")
                                             .RandomInput(type, data_dims)
                                             .Input(indices)
                                             .Attr("T", type)
                                             .Attr("keep_dims", keep_dims));
  });
}

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

TEST_F(OpTest, MaxPool) {
  Repeatedly([this]() {
    std::uniform_int_distribution<int> random_int(1, 5);
    std::vector<int64_t> dims = RandomDims(4, 4, 1);
    int kernel_rows =
        std::uniform_int_distribution<int>(1, dims[1])(generator());
    int kernel_cols =
        std::uniform_int_distribution<int>(1, dims[2])(generator());
    int stride_rows = random_int(generator()),
        stride_cols = random_int(generator());

    std::string padding = Choose<std::string>({"SAME", "VALID"});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("MaxPool")
            .RandomInput(DT_FLOAT, dims)
            .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) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    std::uniform_int_distribution<int> random_int(1, 5);
    std::vector<int64_t> dims = RandomDims(5, 5, 1);

    std::vector<int64_t> input_dims, kernel_dims, stride_dims;
    kernel_dims.push_back(1);
    stride_dims.push_back(1);
    for (int i = 0; i < 3; ++i) {
      kernel_dims.push_back(
          std::uniform_int_distribution<int>(1, dims[i])(generator()));
      input_dims.push_back(dims[i]);
      stride_dims.push_back(random_int(generator()));
    }
    kernel_dims.push_back(1);
    stride_dims.push_back(1);
    int64_t batch = dims[3];
    int64_t feature = dims[4];

    std::string padding = Choose<std::string>({"SAME", "VALID"});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("MaxPool3D")
            .RandomInput(DT_FLOAT,
                         ImageDims(FORMAT_NHWC, batch, feature, input_dims))
            .Attr("T", DT_FLOAT)
            .Attr("ksize", kernel_dims)
            .Attr("strides", stride_dims)
            .Attr("padding", padding)
            .Attr("data_format", "NDHWC"));
  });
  // TODO(phawkins): test NCHW format (not supported by CPU)
}

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

TEST_F(OpTest, Min) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT});
    std::vector<int64_t> data_dims = RandomDims();
    Tensor indices = RandomReductionIndices(data_dims.size());
    bool keep_dims = Choose<bool>({false, true});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Min")
                                             .RandomInput(type, data_dims)
                                             .Input(indices)
                                             .Attr("T", type)
                                             .Attr("keep_dims", keep_dims));
  });
}

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

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

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

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

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

TEST_F(OpTest, NextAfter) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT});
    auto dims = RandomDims();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("NextAfter")
                                             .RandomInput(type, dims)
                                             .RandomInput(type, dims)
                                             .Attr("T", type));
  });
}

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

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

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

    int32_t depth = RandomDim();

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

    int axis = std::uniform_int_distribution<int32_t>(
        -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_t>(depth));
    builder.RandomInput(type, {});
    builder.RandomInput(type, {});
    return ExpectTfAndXlaOutputsAreClose(builder);
  });
}

TEST_F(OpTest, OnesLike) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("OnesLike").RandomInput(type).Attr("T", type));
  });
}

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

    std::vector<int64_t> dims = RandomDims();
    int num_dims = dims.size();
    int axis = std::uniform_int_distribution<int32_t>(-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.RandomInput(type, dims);
    }
    return ExpectTfAndXlaOutputsAreClose(builder);
  });
}

TEST_F(OpTest, Pad) {
  // See note about failing Kokoro tests: b/214080339#comment22
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  Repeatedly([this]() {
    auto a = ChoosePadArguments();
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Pad")
            .RandomInput(a.input_type, a.input_shape)
            .Input(a.paddings)
            .Attr("T", a.input_type)
            .Attr("Tpaddings", a.paddings_type));
  });
}

TEST_F(OpTest, PadV2) {
  Repeatedly([this]() {
    auto a = ChoosePadArguments();
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("PadV2")
            .RandomInput(a.input_type, a.input_shape)
            .Input(a.paddings)
            .Input(a.constant_values)
            .Attr("T", a.input_type)
            .Attr("Tpaddings", a.paddings_type));
  });
}

TEST_F(OpTest, Pow) {
  // TODO(phawkins): Feeding large DT_INT32 values to Pow() leads to
  // nontermination.
  Repeatedly([this]() {
    auto dims = BroadcastableDims();
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Pow")
                                             .RandomInput(type, dims.first)
                                             .RandomInput(type, dims.second)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, Prod) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    std::vector<int64_t> data_dims = RandomDims();
    Tensor indices = RandomReductionIndices(data_dims.size());
    bool keep_dims = Choose<bool>({false, true});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Prod")
                                             .RandomInput(type, data_dims)
                                             .Input(indices)
                                             .Attr("T", type)
                                             .Attr("keep_dims", keep_dims));
  });
}

TEST_F(OpTest, Qr) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Qr")
            .RandomInput(type, RandomDims(2, kDefaultMaxRank, 1))
            .Attr("T", type)
            .Attr("full_matrices", true));
  });
}

TEST_F(OpTest, QuantizeAndDequantizeV2) {
  Repeatedly([this]() {
    std::uniform_int_distribution<int64_t> num_bits_dist(1, 64);
    int64_t num_bits = num_bits_dist(generator());
    std::string round_mode = Choose<std::string>({"HALF_TO_EVEN", "HALF_UP"});
    auto dims = RandomDims(0, kDefaultMaxRank, 1);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("QuantizeAndDequantizeV2")
            .RandomInput(DT_FLOAT, dims)
            .RandomInput(DT_FLOAT, dims)  // unused because range_given = false
            .RandomInput(DT_FLOAT, dims)  // unused because range_given = false
            .Attr("signed_input", RandomBool())
            .Attr("num_bits", num_bits)
            .Attr("range_given", false)
            .Attr("round_mode", round_mode)
            .Attr("narrow_range", RandomBool())
            .Attr("axis", -1)
            .Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, RandomShuffle) {
  // See b/209062491 as this test passes with --tf_xla_test_device=CPU:0
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {  // NOLINT: due to GTEST_SKIP
    auto type = Choose<DataType>(kAllXlaTypes);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("RandomShuffle")
                                             .RandomInput(type, RandomDims(1))
                                             .Attr("seed", RandomSeed())
                                             .Attr("seed2", RandomSeed())
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, RandomStandardNormal) {
  Repeatedly([this]() {
    auto shape_type = Choose<DataType>({DT_INT32, DT_INT64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("RandomStandardNormal")
            .Input(AsIntTensor(shape_type, RandomDims()))
            .Attr("seed", RandomSeed())
            .Attr("seed2", RandomSeed())
            .Attr("T", shape_type)
            .Attr("dtype", DT_FLOAT));
  });
}

TEST_F(OpTest, RandomUniform) {
  Repeatedly([this]() {
    auto shape_type = Choose<DataType>({DT_INT32, DT_INT64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("RandomStandardNormal")
            .Input(AsIntTensor(shape_type, RandomDims()))
            .Attr("seed", RandomSeed())
            .Attr("seed2", RandomSeed())
            .Attr("T", shape_type)
            .Attr("dtype", DT_FLOAT));
  });
}

TEST_F(OpTest, Range) {
  Repeatedly([this]() {
    auto ToScalar = [](DataType type, int x) {
      if (type == DT_INT32) return test::AsScalar<int32_t>(x);
      if (type == DT_INT64) return test::AsScalar<int64_t>(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});
    return 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]() {
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Rank").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, Real) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Real")
                                             .RandomInput(DT_COMPLEX64)
                                             .Attr("T", DT_COMPLEX64));
  });
}

TEST_F(OpTest, RealDiv) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    auto dims = BroadcastableDims();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("RealDiv")
                                             .RandomInput(type, dims.first)
                                             .RandomInput(type, dims.second)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, Reciprocal) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Reciprocal").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, ReciprocalGrad) {
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims();
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ReciprocalGrad")
                                             .RandomInput(type, dims)
                                             .RandomInput(type, dims)
                                             .Attr("T", type));
  });
}
TEST_F(OpTest, Relu) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Relu").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

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

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

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

TEST_F(OpTest, Reshape) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    std::vector<int64_t> dims = RandomDims();
    std::bernoulli_distribution random_bool;
    std::vector<int64_t> dims_before, dims_after;
    for (std::vector<int64_t>* out : {&dims_before, &dims_after}) {
      std::shuffle(dims.begin(), dims.end(), generator());
      for (int64_t 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;
        }
      }
    }
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Reshape")
            .RandomInput(type, dims_before)
            .Input(test::AsTensor<int32_t>(
                std::vector<int32_t>(dims_after.begin(), dims_after.end())))
            .Attr("T", type));
  });
}

TEST_F(OpTest, ResizeNearestNeighbor) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_INT32, DT_INT64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("ResizeNearestNeighbor")
            .RandomInput(type, RandomDims(4, 4, 1))
            .Input(AsIntTensor(DT_INT32, RandomDims(2, kDefaultMaxRank, 1)))
            .Attr("align_corners", RandomBool())
            .Attr("half_pixel_centers", RandomBool())
            .Attr("T", type));
  });
}

TEST_F(OpTest, ResizeBilinear) {
  Repeatedly([this]() {
    std::vector<int64_t> in_dims = RandomDims(4, 4);
    std::vector<int64_t> out_dims = RandomDims(2, 2);

    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("ResizeBilinear")
            .RandomInput(DT_FLOAT, in_dims)
            .Input(test::AsTensor<int32_t>(
                std::vector<int32_t>(out_dims.begin(), out_dims.end())))
            .Attr("T", DT_FLOAT)
            .Attr("align_corners", true));
  });
}

TEST_F(OpTest, ResizeBilinearGrad) {
  Repeatedly([this]() {
    std::vector<int64_t> in_dims = RandomDims(4, 4);
    std::vector<int64_t> out_dims = RandomDims(2, 2);

    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("ResizeBilinearGrad")
            .RandomInput(DT_FLOAT, in_dims)
            .RandomInput(DT_FLOAT,
                         {in_dims[0], out_dims[0], out_dims[1], in_dims[3]})
            .Attr("T", DT_FLOAT)
            .Attr("align_corners", true));
  });
}

TEST_F(OpTest, Reverse) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims(1);
    auto type = Choose<DataType>(kAllXlaTypes);
    int64_t rank = dims.size();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Reverse")
                                             .RandomInput(type, dims)
                                             .RandomInput(DT_BOOL, {rank})
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, ReverseSequence) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims(/*min_rank=*/2);
    auto type = Choose<DataType>(kAllXlaTypes);
    int64_t rank = dims.size();

    // Choose random batch and sequence dimensions.
    std::vector<int> shuffled_dim_ids(rank);
    absl::c_iota(shuffled_dim_ids, 0);
    absl::c_shuffle(shuffled_dim_ids, generator());
    shuffled_dim_ids.resize(2);
    int batch_dim = shuffled_dim_ids[0];
    int seq_dim = shuffled_dim_ids[1];

    int batch_size = dims[batch_dim];
    int max_seq_len = dims[seq_dim];
    std::vector<int32_t> seq_lens(batch_size);
    std::uniform_int_distribution<int32_t> d(0, max_seq_len);
    absl::c_generate(seq_lens, [&]() { return d(generator()); });

    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("ReverseSequence")
            .RandomInput(type, dims)
            .Input(test::AsTensor<int32_t>(seq_lens))
            .Attr("seq_dim", seq_dim)
            .Attr("batch_dim", batch_dim)
            .Attr("T", type)
            .Attr("Tlen", DT_INT32));
  });
}

TEST_F(OpTest, ReverseV2) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    std::vector<int64_t> data_dims = RandomDims();
    Tensor indices = RandomReductionIndices(data_dims.size());
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ReverseV2")
                                             .RandomInput(type, data_dims)
                                             .Input(indices)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, RightShift) {
  Repeatedly([this]() {
    bool is64 = RandomBool();
    auto dims = RandomDims();
    auto type = is64 ? DT_INT64 : DT_INT32;
    int max_shift = is64 ? 63 : 31;
    auto y = RandomBoundedTensor(type, 0, max_shift, false, dims);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("RightShift")
                                             .RandomInput(type, dims)
                                             .Input(y)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, Rint) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Rint").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, Roll) {
  Repeatedly([this]() {
    auto input_type = Choose<DataType>(kAllXlaTypes);
    auto axis_type = Choose<DataType>({DT_INT32, DT_INT64});
    // TODO(b/201095155,b/197140886): shift_type = DT_INT64 doesn't work.
    auto shift_type = DT_INT32;
    auto input_shape = RandomDims(1);
    int rank = input_shape.size();
    auto axis_shape = RandomDims(1, 1, 1, rank + 1);
    auto axis = RandomBoundedTensor(axis_type, 0, rank - 1, true, axis_shape);
    auto shift = RandomTensor(shift_type, false, axis_shape);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Roll")
            .RandomInput(input_type, input_shape)
            .Input(shift)
            .Input(axis)
            .Attr("T", input_type)
            .Attr("Taxis", axis_type)
            .Attr("Tshift", shift_type));
  });
}

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

TEST_F(OpTest, Rsqrt) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Rsqrt").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, RsqrtGrad) {
  Repeatedly([this]() {
    auto dims = RandomDims();
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("RsqrtGrad")
                                             .RandomInput(type, dims)
                                             .RandomInput(type, dims)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, Select) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    auto shape = RandomDims();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Select")
                                             .RandomInput(DT_BOOL, shape)
                                             .RandomInput(type, shape)
                                             .RandomInput(type, shape)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, SelectV2) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    auto shape = RandomDims();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SelectV2")
                                             .RandomInput(DT_BOOL, shape)
                                             .RandomInput(type, shape)
                                             .RandomInput(type, shape)
                                             .Attr("T", type));
  });
}

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

TEST_F(OpTest, ShapeN) {
  Repeatedly([this]() {
    auto 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.RandomInput(type);
    }
    return ExpectTfAndXlaOutputsAreClose(builder);
  });
}

TEST_F(OpTest, Sigmoid) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Sigmoid").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, SigmoidGrad) {
  Repeatedly([this]() {
    auto dims = RandomDims();
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SigmoidGrad")
                                             .RandomInput(type, dims)
                                             .RandomInput(type, dims)
                                             .Attr("T", type));
  });
}

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

TEST_F(OpTest, Sin) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Sin").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, Sinh) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Sinh").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, Size) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Size").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, Slice) {
  Repeatedly([this]() {
    SliceArguments a = ChooseSliceArguments(true);
    std::vector<int32_t> size;
    size.insert(size.end(), a.size.begin(), a.size.end());
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Slice")
            .RandomInput(a.type, a.shape)
            .Input(a.indices)
            .Input(test::AsTensor<int32_t>(size))
            .Attr("T", a.type)
            .Attr("Index", a.indices_type));
  });
}

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

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

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

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

TEST_F(OpTest, Softsign) {
  Repeatedly([this]() {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Softsign").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT));
  });
}

TEST_F(OpTest, SoftsignGrad) {
  Repeatedly([this]() {
    std::vector<int64_t> dims = RandomDims();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SoftsignGrad")
                                             .RandomInput(DT_FLOAT, dims)
                                             .RandomInput(DT_FLOAT, dims)
                                             .Attr("T", DT_FLOAT));
  });
}

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

    std::vector<int64_t> 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_t> padding_vals;
    std::uniform_int_distribution<int> distribution(0, 7);
    for (int i = 0; i < num_block_dims; ++i) {
      int64_t pad_before;
      int64_t 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})));

    auto type = Choose<DataType>(kAllXlaTypes);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SpaceToBatch")
                                             .RandomInput(type, input_dims)
                                             .Input(paddings)
                                             .Attr("T", type)
                                             .Attr("block_size", block_size));
  });
}

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

    std::vector<int64_t> 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_t> padding_vals;
    std::uniform_int_distribution<int> distribution(0, 7);
    for (int i = 0; i < num_block_dims; ++i) {
      int64_t pad_before;
      int64_t 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})));

    auto type = Choose<DataType>(kAllXlaTypes);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("SpaceToBatchND")
            .RandomInput(type, input_dims)
            .Input(test::AsTensor<int32_t>(
                std::vector<int32_t>(block_dims.begin(), block_dims.end())))
            .Input(paddings)
            .Attr("T", type));
  });
}

TEST_F(OpTest, SpaceToDepth) {
  Repeatedly([this]() {
    int64_t block = RandomDim(2, 5);
    std::vector<int64_t> input_dims = RandomDims(4, 4);
    // Round spatial dimensions up to a multiple of the block size
    input_dims[1] = (input_dims[1] + (block - 1)) / block * block;
    input_dims[2] = (input_dims[2] + (block - 1)) / block * block;
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SpaceToDepth")
                                             .RandomInput(DT_FLOAT, input_dims)
                                             .Attr("T", DT_FLOAT)
                                             .Attr("block_size", block));
  });
}

TEST_F(OpTest, SparseMatMul) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    int64_t x = RandomDim();
    int64_t y = RandomDim();
    int64_t z = RandomDim();

    std::vector<int64_t> a_dims = {x, y};
    std::vector<int64_t> b_dims = {y, z};

    std::bernoulli_distribution random_bool;
    bool transpose_a = random_bool(generator());
    bool transpose_b = random_bool(generator());
    if (transpose_a) {
      std::swap(a_dims[0], a_dims[1]);
    }
    if (transpose_b) {
      std::swap(b_dims[0], b_dims[1]);
    }

    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SparseMatMul")
                                             .RandomInput(DT_FLOAT, a_dims)
                                             .RandomInput(DT_FLOAT, b_dims)
                                             .Attr("Ta", DT_FLOAT)
                                             .Attr("Tb", DT_FLOAT)
                                             .Attr("transpose_a", transpose_a)
                                             .Attr("transpose_b", transpose_b));
  });
}

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

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

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

TEST_F(OpTest, Split) {
  // See b/214080339#comment27 as this test causes Kokoro to crash.
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    std::vector<int64_t> dims = RandomDims(1);
    std::uniform_int_distribution<int> ud;
    int32_t dim = std::uniform_int_distribution<int32_t>(
        -static_cast<int32_t>(dims.size()),
        static_cast<int32_t>(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;
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Split")
            .Input(test::AsScalar<int32_t>(dim))
            .RandomInput(type, dims)
            .Attr("T", type)
            .Attr("num_split", n));
  });
}

TEST_F(OpTest, SplitV) {
  // Likely this only fails when dim is negative. Try type = DT_FLOAT first.
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {  // NOLINT: due to GTEST_SKIP
    auto type = Choose<DataType>(kAllXlaTypes);
    std::vector<int64_t> dims = RandomDims(1, kDefaultMaxRank, 1);
    int32_t dim = std::uniform_int_distribution<int32_t>(
        -static_cast<int32_t>(dims.size()),
        static_cast<int32_t>(dims.size()) - 1)(generator());
    int n = std::uniform_int_distribution<int>(
        1, std::min(5, static_cast<int>(dims[dim])))(generator());
    std::vector<int32_t> size_splits(n);
    for (int i = 0; i < n - 1; ++i) {
      size_splits.push_back(dims[dim] / n);
    }
    size_splits.push_back(dims[dim] - (n - 1) * (dims[dim] / n));
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("SplitV")
            .RandomInput(type, dims)
            .Input(test::AsTensor<int32_t>(size_splits))
            .Input(test::AsScalar<int32_t>(dim))
            .Attr("T", type)
            .Attr("num_split", n)
            .Attr("Tlen", DT_INT32));
  });
}

TEST_F(OpTest, Sqrt) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Sqrt").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, StopGradient) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("StopGradient").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, SqrtGrad) {
  Repeatedly([this]() {
    auto dims = RandomDims();
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SqrtGrad")
                                             .RandomInput(type, dims)
                                             .RandomInput(type, dims)
                                             .Attr("T", type));
  });
}

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

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

TEST_F(OpTest, Squeeze) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    std::vector<int64_t> t_dims = RandomDims(0, kDefaultMaxRank, 0, 5);
    std::bernoulli_distribution random_bool;
    std::vector<int> squeeze_dims;
    for (int i = 0; i < t_dims.size(); ++i) {
      if (t_dims[i] == 1 && random_bool(generator())) {
        squeeze_dims.push_back(i);
      }
    }
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Squeeze")
                                             .RandomInput(type, t_dims)
                                             .Attr("squeeze_dims", squeeze_dims)
                                             .Attr("T", type));
  });
}

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

TEST_F(OpTest, Sum) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    std::vector<int64_t> data_dims = RandomDims();
    Tensor indices = RandomReductionIndices(data_dims.size());
    bool keep_dims = Choose<bool>({false, true});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Sum")
                                             .RandomInput(type, data_dims)
                                             .Input(indices)
                                             .Attr("T", type)
                                             .Attr("keep_dims", keep_dims));
  });
}

TEST_F(OpTest, StridedSlice) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    std::vector<int64_t> data_dims = RandomDims();
    std::vector<int32_t> begin(data_dims.size()), end(data_dims.size());
    std::vector<int32_t> strides(data_dims.size());
    for (int i = 0; i < data_dims.size(); ++i) {
      begin[i] = std::uniform_int_distribution<int32_t>(
          -2 * data_dims[i], 2 * data_dims[i])(generator());
      end[i] = std::uniform_int_distribution<int32_t>(
          -2 * data_dims[i], 2 * data_dims[i])(generator());
      // TODO(b/31360685): support strides other than 1 or -1
      strides[i] = std::bernoulli_distribution()(generator()) ? 1 : -1;
    }
    int64_t max_bitmask = (1LL << data_dims.size()) - 1;
    std::uniform_int_distribution<int64_t> bitmask_distribution(0, max_bitmask);
    int64_t begin_mask = bitmask_distribution(generator());
    int64_t end_mask = bitmask_distribution(generator());

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

    int64_t new_axis_mask = bitmask_distribution(generator());
    int64_t shrink_axis_mask = bitmask_distribution(generator());
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("StridedSlice")
            .RandomInput(type, data_dims)
            .Input(test::AsTensor<int32_t>(begin))
            .Input(test::AsTensor<int32_t>(end))
            .Input(test::AsTensor<int32_t>(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) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);

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

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

    // Create a ellipsis bitmask with at most one 1 bit set.
    int64_t 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_t new_axis_mask = bitmask_distribution(generator());
    int64_t 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.
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("StridedSliceGrad")
            .Input(test::AsTensor<int64_t>(dims))
            .Input(test::AsTensor<int64_t>(begin))
            .Input(test::AsTensor<int64_t>(end))
            .Input(test::AsTensor<int64_t>(strides))
            .RandomInput(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, Tan) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Tan").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, Tanh) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Tanh").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, TanhGrad) {
  Repeatedly([this]() {
    auto dims = RandomDims();
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("TanhGrad")
                                             .RandomInput(type, dims)
                                             .RandomInput(type, dims)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, TensorScatterUpdate) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {  // NOLINT: due to GTEST_SKIP
    auto a = ChooseScatterArguments();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("TensorScatterUpdate")
                                             .RandomInput(a.type, a.shape)
                                             .Input(a.indices)
                                             .Input(a.updates)
                                             .Attr("T", a.type)
                                             .Attr("Tindices", a.indices_type));
  });
}

TEST_F(OpTest, Tile) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    std::vector<int64_t> t_dims = RandomDims(1);
    std::vector<int32_t> multiples(t_dims.size());
    for (int i = 0; i < t_dims.size(); ++i) {
      multiples[i] = std::uniform_int_distribution<int>(1, 3)(generator());
    }
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Tile")
            .RandomInput(type, t_dims)
            .Input(test::AsTensor<int32_t>(multiples))
            .Attr("T", type));
  });
}

TEST_F(OpTest, TopKV2) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {  // NOLINT: due to GTEST_SKIP
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_INT64});
    auto shape = RandomDims(1);
    int32_t k =
        std::uniform_int_distribution<int32_t>(1, shape[0])(generator());
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("TopKV2")
                                             .RandomInput(type, shape)
                                             .Input(test::AsScalar<int32_t>(k))
                                             .Attr("sorted", RandomBool())
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, Transpose) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    std::vector<int64_t> data_dims = RandomDims();
    std::vector<int32_t> perm(data_dims.size());
    std::iota(perm.begin(), perm.end(), 0);
    std::shuffle(perm.begin(), perm.end(), generator());
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("Transpose")
            .RandomInput(type, data_dims)
            .Input(test::AsTensor<int32_t>(perm))
            .Attr("T", type));
  });
}

TEST_F(OpTest, TruncateDiv) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  Repeatedly([this]() {
    DataType type = DT_INT32;
    auto dims = BroadcastableDims();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("TruncateDiv")
                                             .RandomInput(type, dims.first)
                                             .RandomInput(type, dims.second)
                                             .Attr("T", type));
  });
}

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

TEST_F(OpTest, Unpack) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    auto shape = RandomDims(1);
    int axis =
        std::uniform_int_distribution<int>(0, shape.size() - 1)(generator());
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Unpack")
                                             .RandomInput(type, shape)
                                             .Attr("axis", axis)
                                             .Attr("T", type)
                                             .Attr("num", shape[axis]));
  });
}

TEST_F(OpTest, Xdivy) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    auto dims = BroadcastableDims();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Xdivy")
                                             .RandomInput(type, dims.first)
                                             .RandomInput(type, dims.second)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, XlaDot) {
  Repeatedly([this]() {
    const XlaDotArguments& a = ChooseXlaDotArguments();
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("XlaDot")
            .RandomInput(a.dtype, a.lhs_dims)
            .RandomInput(a.dtype, a.rhs_dims)
            .Attr("dimension_numbers", a.dnums_encoded)
            .Attr("precision_config", a.precision_config_encoded)
            .Attr("T", a.dtype));
  });
}

TEST_F(OpTest, XlaDotV2) {
  Repeatedly([this]() {
    const XlaDotArguments& a = ChooseXlaDotArguments();
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("XlaDotV2")
            .RandomInput(a.dtype, a.lhs_dims)
            .RandomInput(a.dtype, a.rhs_dims)
            .Attr("dimension_numbers", a.dnums_encoded)
            .Attr("precision_config", a.precision_config_encoded)
            .Attr("LhsT", a.dtype)
            .Attr("RhsT", a.dtype)
            .Attr("preferred_element_type", a.dtype));
  });
}

TEST_F(OpTest, XlaDynamicUpdateSlice) {
  Repeatedly([this]() {
    SliceArguments a = ChooseSliceArguments(false);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("XlaDynamicUpdateSlice")
                                             .RandomInput(a.type, a.shape)
                                             .RandomInput(a.type, a.size)
                                             .Input(a.indices)
                                             .Attr("T", a.type)
                                             .Attr("Tindices", a.indices_type));
  });
}

TEST_F(OpTest, XlaEinsum) {
  Repeatedly([this]() {
    const EinsumArguments a = ChooseEinsumArguments();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("XlaEinsum")
                                             .RandomInput(a.type, a.lhs_dims)
                                             .RandomInput(a.type, a.rhs_dims)
                                             .Attr("equation", a.equation)
                                             .Attr("T", a.type));
  });
}

TEST_F(OpTest, XlaSort) {
  Repeatedly([this]() {
    auto type = Choose<DataType>(kAllXlaTypes);
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("XlaSort")
                                             .RandomInput(type, RandomDims())
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, Xlog1py) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    auto dims = BroadcastableDims();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Xlog1py")
                                             .RandomInput(type, dims.first)
                                             .RandomInput(type, dims.second)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, Xlogy) {
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_FLOAT, DT_COMPLEX64});
    auto dims = BroadcastableDims();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Xlogy")
                                             .RandomInput(type, dims.first)
                                             .RandomInput(type, dims.second)
                                             .Attr("T", type));
  });
}

TEST_F(OpTest, ZerosLike) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  Repeatedly([this]() {
    auto type = Choose<DataType>({DT_INT32, DT_FLOAT, DT_COMPLEX64});
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("ZerosLike").RandomInput(type).Attr("T", type));
  });
}

TEST_F(OpTest, Zeta) {
  Repeatedly([this]() {
    auto dims = BroadcastableDims();
    return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Xlogy")
                                             .RandomInput(DT_FLOAT, dims.first)
                                             .RandomInput(DT_FLOAT, dims.second)
                                             .Attr("T", DT_FLOAT));
  });
}

// Example failing run:
//   --tf_xla_reference_device=GPU:0
//   --tf_xla_test_use_jit=true --tf_xla_test_device=GPU:0
//   --tf_xla_test_use_mlir=true
//   --tf_xla_test_repetitions=2
//   --gunit_filter='OpTest.FusedBatchNormTraining'
//   --tf_xla_random_seed=2838146746
TEST_F(OpTest, FusedBatchNormTraining) {
  if (tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/201095155";
  if (!tensorflow::tf_xla_test_use_mlir) GTEST_SKIP() << "b/197140886";
  bool is_nhwc = RandomBool();
  std::vector<int64_t> x_dims = RandomDims(/*min_rank=*/4, /*max_rank=*/4,
                                           /*min_size=*/5, /*max_size=*/20);
  std::vector<int64_t> scale_dims = {x_dims[is_nhwc ? 3 : 1]};
  std::vector<int64_t> offset_dims = {x_dims[is_nhwc ? 3 : 1]};
  std::vector<int64_t> mean_dims = {0};
  std::vector<int64_t> variance_dims = {0};
  DataType type = DT_FLOAT;
  Repeatedly([&] {
    return ExpectTfAndXlaOutputsAreClose(
        OpTestBuilder("FusedBatchNorm")
            .RandomInput(type, x_dims)
            .RandomInput(type, scale_dims)
            .RandomInput(type, offset_dims)
            .RandomInput(type, mean_dims)
            .RandomInput(type, variance_dims)
            .Attr("T", type)
            .Attr("data_format", is_nhwc ? "NHWC" : "NCHW")
            .Attr("epsilon", static_cast<float>(1.001e-05))
            .Attr("is_training", true));
  });
}
}  // anonymous namespace
}  // namespace tensorflow

int main(int argc, char** argv) {
  tensorflow::tf_xla_test_device_ptr = new std::string("GPU:0");
  tensorflow::tf_xla_reference_device_ptr = new std::string("CPU: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 "
          "nondeterministically."),
      // 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_max_tensor_size",
                       &tensorflow::tf_xla_max_tensor_size,
                       "Maximum number of elements for random input tensors."),
      tensorflow::Flag("tf_xla_test_device", tensorflow::tf_xla_test_device_ptr,
                       "Tensorflow device type to use for test"),
      tensorflow::Flag("tf_xla_reference_device",
                       tensorflow::tf_xla_reference_device_ptr,
                       "Tensorflow device type to use for reference"),
      tensorflow::Flag("tf_xla_test_use_jit", &tensorflow::tf_xla_test_use_jit,
                       "Use JIT compilation for the operator under test"),
      tensorflow::Flag(
          "tf_xla_test_use_mlir", &tensorflow::tf_xla_test_use_mlir,
          "Use MLIR legalization kernels for the operator under test"),
  };
  std::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<std::unique_ptr<tensorflow::Device>> devices;
  CHECK_OK(tensorflow::DeviceFactory::AddDevices(tensorflow::SessionOptions(),
                                                 "", &devices));
  tensorflow::StaticDeviceMgr device_mgr(std::move(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)?";

  if (tensorflow::tf_xla_test_use_mlir)
    tensorflow::GetMlirCommonFlags()->tf_mlir_enable_mlir_bridge =
        tensorflow::ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_ENABLED;
  return RUN_ALL_TESTS();
}