Function IsEqualComputation

tensorflow/compiler/xrt/tests/raw_api_test.cc:257–267

}

xla::XlaComputation IsEqualComputation(const xla::Shape& shape) {
  xla::XlaBuilder builder("IsEqualComputation");
  auto p0 = xla::Parameter(&builder, 0, shape, "P0");
  auto p1 = xla::Parameter(&builder, 1, shape, "P1");
  auto cmp =
      xla::Ne(xla::Sub(p0, p1), xla::Zero(&builder, shape.element_type()));
  auto icmp = xla::ConvertElementType(cmp, xla::S32);
  xla::ReduceAll(icmp, xla::Zero(&builder, xla::S32),
                 xla::CreateScalarAddComputation(xla::S32, &builder));
  return builder.Build().ValueOrDie();
}

void StoreComputationSnapshot(const xla::XlaComputation& computation,
                              xla::HloSnapshot* dst) {