int64_t device_ordinal0 = -1;
  mlir::LogicalResult result0 =
      GetDeviceOrdinalFromDeviceString(unknown_loc, tpu0, &device_ordinal0);
  EXPECT_TRUE(mlir::succeeded(result0));
  EXPECT_EQ(device_ordinal0, 0);

  int64_t device_ordinal1 = -1;
  mlir::LogicalResult result1 =
      GetDeviceOrdinalFromDeviceString(unknown_loc, tpu1, &device_ordinal1);
  EXPECT_TRUE(mlir::succeeded(result1));

  XlaCompiler::Options options;
  options.device_ordinal = 0;
  options.alias_passthrough_params = true;
  options.detailed_logging = true;
  XlaCompiler::CompilationResult result;
  xla::Shape xla_output_shape;
  result.xla_output_shape = xla_output_shape;

  auto build_option =
      GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/-1);

  EXPECT_EQ(build_option.device_ordinal(), 0);

inline void LogOptions(const XlaCompiler::Options& options) {
  VLOG(2) << "XlaCompiler::Options[device_type=" << options.device_type
          << ",device_ordinal=" << options.device_ordinal
          << ",client=" << options.client << ",flib_def=" << options.flib_def
          << ",graph_def_version=" << options.graph_def_version
          << ",options.shape_determination_fns.layout_preference_fn?="

  // CHECK-SAME:   device_ordinal = 0
  // CHECK:      "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"
  // CHECK-SAME:   device_ordinal = 1
  // CHECK:      "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"
  // CHECK-SAME:   device_ordinal = 0
  // CHECK:      "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"
  // CHECK-SAME:   device_ordinal = 1
}

// -----

      if (op->getAttrOfType<IntegerAttr>(kDeviceOrdinalAttr).getInt() == -1) {
        OpBuilder builder(op);
        int device_ordinal =
            op->getAttrOfType<IntegerAttr>(kReplicaIdAttr).getInt() %
            device_num;
        op->setAttr(kDeviceOrdinalAttr,
                    builder.getI64IntegerAttr(device_ordinal));
      }
    }
  }
}
}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>>

  "tf.EnqueueTPUEmbeddingSparseTensorBatch"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %0, %0, %0, %arg7) {_tpu_embedding_layer = "call1", _xla_outside_compilation = "0", combiners = ["mean", "sum"], device_ordinal = -1 : i64, max_sequence_lengths = [0, 0, 0], table_ids = [1, 1, 0]} : (tensor<?x2xi32>, tensor<?x2xi32>, tensor<?x2xi32>, tensor<?xi32>, tensor<?xi32>, tensor<?xi32>, tensor<0xf32>, tensor<0xf32>, tensor<0xf32>, tensor<!tf_type.string>) -> ()...

    Options(std::string device_name_prefix, std::string device_name,
            int device_ordinal, std::string compilation_device_name,
            std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
                shape_determination_fns)
        : device_name_prefix(device_name_prefix),
          device_name(device_name),
          device_ordinal(device_ordinal),
          compilation_device_name(compilation_device_name),

                                               const string& device_name,
                                               int device_ordinal) {
  return Device::BuildDeviceAttributes(
      absl::StrCat(name_prefix, "/device:", device_name, ":", device_ordinal),
      DeviceType(device_name), Bytes(16ULL << 30), DeviceLocality(),
      absl::StrCat("device: ", device_name, " device"));
}

}  // namespace

  xla::ExecutableBuildOptions build_options;
  if (result.collective_info) {
    build_options.set_num_replicas(result.collective_info->group_size);
  }
  if (options.device_ordinal != -1) {
    build_options.set_device_ordinal(options.device_ordinal);
  } else if (default_device_ordinal != -1) {
    build_options.set_device_ordinal(default_device_ordinal);
  }
  build_options.set_result_layout(result.xla_output_shape);

                                       int device_ordinal) {
  xla::Shape on_host_shape =
      xla::ShapeUtil::DeviceShapeToHostShape(on_device_shape);
  xla::ScopedShapedBuffer shaped_buffer(on_host_shape, on_device_shape,
                                        client->backend().memory_allocator(),
                                        device_ordinal);
  for (auto& index_to_buffer : shaped_buffer.buffers()) {