Function TestFusedAdamBase

test/cpp/phi/kernels/test_fused_adam_kernel.cc:428–475

template <typename T, typename PlaceType>
void TestFusedAdamBase(const std::vector<std::vector<int64_t>> &shapes,
                       float atol,
                       bool use_adamw,
                       bool amsgrad,
                       bool multi_precision = false,
                       float beta1 = 0.9,
                       float beta2 = 0.99,
                       float weight_decay = 0.1,
                       size_t steps = 5,
                       uint64_t seed = 10) {
  const auto &ctx = *phi::DeviceContextPool::Instance().GetByPlace(PlaceType());
  using Context = typename std::remove_const<
      typename std::remove_pointer<decltype(&ctx)>::type>::type;
  ctx.GetGenerator()->SetCurrentSeed(seed);
  AdamInfo<T, Context> info1(ctx,
                             shapes,
                             beta1,
                             beta2,
                             weight_decay,
                             multi_precision,
                             use_adamw,
                             amsgrad);
  auto info2 = AdamInfo<T, Context>::DeepCopy(info1);

  for (size_t i = 0; i < steps; ++i) {
    auto grads = GenerateRandomTensorVectors<T>(ctx, shapes);
    info1.Update(false, grads);
    info2.Update(true, grads);
  }

  using MT = typename decltype(info1)::MT;

#define PD_ADAM_TEST_COMP(__field, __dtype)                          \
  do {                                                               \
    MT __diff = MaxDiff<__dtype>(ctx, info1.__field, info2.__field); \
    EXPECT_LE(__diff, static_cast<MT>(atol))                         \
        << #__field << " has diff when use_adamw = " << use_adamw    \
        << " , multi_precision = " << multi_precision;               \
  } while (0)

  PD_ADAM_TEST_COMP(beta1_pows, MT);
  PD_ADAM_TEST_COMP(beta2_pows, MT);
  PD_ADAM_TEST_COMP(params, T);
  PD_ADAM_TEST_COMP(master_params, MT);
  PD_ADAM_TEST_COMP(moment1s, MT);
  PD_ADAM_TEST_COMP(moment2s, MT);
  PD_ADAM_TEST_COMP(moment2s_max, MT);
}

static auto GenerateRandomShapes(size_t n, uint64_t low, uint64_t high) {
  std::random_device device;