ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  float product[4] = {0};
  EXPECT_EQ(sizeof(product), TF_TensorByteSize(t));
  memcpy(&product[0], TF_TensorData(t), TF_TensorByteSize(t));
  TF_DeleteTensor(t);
  EXPECT_EQ(7, product[0]);
  EXPECT_EQ(10, product[1]);
  EXPECT_EQ(15, product[2]);
  EXPECT_EQ(22, product[3]);
  TF_DeleteStatus(status);
}

  // persistent cache.
  std::vector<XlaJitCompilationActivity> activity_history =
      listener_->GetListenerHistory();
  EXPECT_EQ(activity_history.size(), 1);
  EXPECT_EQ(activity_history[0].cluster_name(), fn.name());
  EXPECT_EQ(activity_history[0].compile_count(), 1);
  EXPECT_FALSE(activity_history[0].used_persistent_cache());

  listener_->ClearListenerHistory();

                const std::vector<string>& want_namespaces) {
    string class_name;
    std::vector<string> namespaces;
    TF_EXPECT_OK(ParseCppClass(cpp_class, &class_name, &namespaces));
    EXPECT_EQ(class_name, want_class_name);
    EXPECT_EQ(namespaces, want_namespaces);
  }

  void ExpectFail(const string& cpp_class) {
    string class_name;
    std::vector<string> namespaces;

        /*device=*/""}});
}

template <typename T>
void ExpectEquals(AbstractTensorHandle* t, T expected) {
  TF_Tensor* result_t;
  Status s = tensorflow::GetValue(t, &result_t);
  ASSERT_TRUE(s.ok()) << s.message();
  auto value = static_cast<T*>(TF_TensorData(result_t));
  EXPECT_EQ(*value, expected);
  TF_DeleteTensor(result_t);
}

// TODO(srbs): Add tests for captures.

      GetVariableSnapshots(variables),
      /*pjrt_client=*/pjrt_client_, /*pjrt_device=*/*pjrt_device,
      /*use_pjrt_tensor_buffer=*/true, &exec_args,
      /*owned_args=*/{}, &non_donatable_input_indices));

  EXPECT_EQ(exec_args.size(), 2);

  std::shared_ptr<xla::Literal> literal1 = *exec_args[0]->ToLiteralSync();
  EXPECT_TRUE(xla::LiteralTestUtil::Equal(
      *literal1, xla::LiteralUtil::CreateR2<int32_t>({{1, 2, 3}})));

    TF_EXPECT_OK(
        session.Run({{x, {-3.0f, 0.0f, 3.0f}}}, grad_outputs, &grad_result));
    EXPECT_EQ(grad_result.size(), 1);
    EXPECT_EQ(grad_result[0].NumElements(), 3);
    EXPECT_EQ(grad_result[0].flat<float>()(0), 0.0f);
    EXPECT_EQ(grad_result[0].flat<float>()(1), 0.0f);
    EXPECT_EQ(grad_result[0].flat<float>()(2), 0.0f);
  }
}

TEST_F(NaryGradTest, SquaredDifference) {

  root.graph()->AddControlEdge(call, write_op);

  std::unique_ptr<Graph> graph;
  Status failure_status = BuildXlaOps(root, fdef_lib, &graph);
  ASSERT_FALSE(failure_status.ok());
  EXPECT_EQ(failure_status.code(), error::INVALID_ARGUMENT);
}

TEST_F(BuildXlaOpsTest, OnNonXlaDevice) {
  Scope root = Scope::NewRootScope().ExitOnError();

  FunctionDefLibrary fdef_lib =

  std::optional<CalibrationStatistics> statistics = collector.GetStatistics();

  EXPECT_TRUE(statistics.has_value());
  EXPECT_EQ(statistics.value().average_min_max_statistics().min_sum(), -4.0f);
  EXPECT_EQ(statistics.value().average_min_max_statistics().max_sum(), 15.0f);
  EXPECT_EQ(statistics.value().average_min_max_statistics().num_samples(), 2);
}

TEST(CalibrationStatisticsCollectorTest, ClearDataAndGetResultsMinMax) {

            }
        """

        and:
        module.jar.expectPut()
        module.jar.sha1.expectPut()
        module.jar.sha256.expectPut()
        module.jar.sha512.expectPut()
        module.ivy.expectPut()
        module.ivy.sha1.expectPut()
        module.ivy.sha256.expectPut()
        module.ivy.sha512.expectPut()

        when:
        run 'publish'

        then:

  // verify transpose
  EXPECT_EQ(
      fused_lstm_func_->getAttrOfType<StringAttr>(kTFImplements).getValue(),
      convert.GetCompositeOpName());
  EXPECT_EQ(fused_lstm_func_.getNumArguments(), 5);
  EXPECT_EQ(fused_lstm_func_.getFunctionType().getNumResults(), 1);

  auto transpose_op = fused_lstm_func_.getBody().front().begin();
  transpose_op++;