// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include <random>

#include "gtest/gtest.h"
#include "test/providers/provider_test_utils.h"
#include "test/util/include/default_providers.h"

namespace onnxruntime {
namespace test {

namespace {

TEST(OptimizerTest, SGDOptimizerTest) {
  OpTester test("SGDOptimizer", 1, onnxruntime::kMSDomain);
  test.AddInput<float>("ETA", {}, {0.5f});
  test.AddInput<float>("W", {3}, {1, 2, 3});
  test.AddInput<float>("G", {3}, {4, 5, 6});
  test.AddOutput<float>("W_New", {3}, {-1.f, -0.5f, 0.f});
  test.Run();
}

TEST(OptimizerTest, SGDOptimizerTest_Gradient) {
  OpTester test("SGDOptimizer", 1, onnxruntime::kMSDomain);
  test.AddInput<float>("ETA", {}, {0.5f});
  test.AddInput<float>("W", {3}, {1, 2, 3});
  test.AddInput<float>("G", {3}, {4, 5, 6});
  test.AddOptionalOutputEdge<float>();
  test.AddOutput<float>("G_New", {3}, {-2.f, -2.5f, -3.f});
  test.Run();
}

struct AdamOptimizerInputOutput {
  AdamOptimizerInputOutput() {
    eta_half.resize(eta.size());
    g_half.resize(g.size());
    m1_half.resize(m1.size());
    m2_half.resize(m2.size());
    w_half.resize(w.size());
    ConvertFloatToMLFloat16(eta.data(), eta_half.data(), int(eta.size()));
    ConvertFloatToMLFloat16(g.data(), g_half.data(), int(g.size()));
    ConvertFloatToMLFloat16(m1.data(), m1_half.data(), int(m1.size()));
    ConvertFloatToMLFloat16(m2.data(), m2_half.data(), int(m2.size()));
    ConvertFloatToMLFloat16(w.data(), w_half.data(), int(w.size()));

    m1_new_half.resize(m1_new.size());
    m2_new_half.resize(m2_new.size());
    w_new_half.resize(w_new.size());
    g_new_half.resize(g_new.size());
    ConvertFloatToMLFloat16(m1_new.data(), m1_new_half.data(), int(m1_new.size()));
    ConvertFloatToMLFloat16(m2_new.data(), m2_new_half.data(), int(m2_new.size()));
    ConvertFloatToMLFloat16(w_new.data(), w_new_half.data(), int(w_new.size()));
    ConvertFloatToMLFloat16(g_new.data(), g_new_half.data(), int(g_new.size()));
  }

  // Fp32 Inputs
  std::vector<float> eta = {0.5f};
  std::vector<float> w = {1.0f, 2.0f, 3.0f};
  std::vector<float> g = {4.0f, 5.0f, 6.0f};
  std::vector<float> m1 = {0.1f, 0.2f, 0.3f};
  std::vector<float> m2 = {0.4f, 0.5f, 0.6f};

  // Fp16 Inputs
  std::vector<MLFloat16> eta_half;
  std::vector<MLFloat16> w_half;
  std::vector<MLFloat16> g_half;
  std::vector<MLFloat16> m1_half;
  std::vector<MLFloat16> m2_half;

  // FP32 Outptus
  std::vector<float> m1_new = {0.49f, 0.68f, 0.87f};
  std::vector<float> m2_new = {0.4156f, 0.5245f, 0.6354f};
  std::vector<float> w_new = {0.6199609f, 1.5305318f, 2.4542853f};
  std::vector<float> g_new = {-0.3800391f, -0.4694682f, -0.5457147f};

  // FP16 Outptus
  std::vector<MLFloat16> m1_new_half;
  std::vector<MLFloat16> m2_new_half;
  std::vector<MLFloat16> w_new_half;
  std::vector<MLFloat16> g_new_half;
};

TEST(OptimizerTest, AdamOptimizerTest) {
  OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
  AdamOptimizerInputOutput data;

  test.AddInput<float>("ETA", {}, data.eta);
  test.AddInput<int64_t>("Update_Count", {}, {3});
  test.AddInput<float>("W", {3}, data.w);
  test.AddInput<float>("G", {3}, data.g);
  test.AddInput<float>("Moment_1", {3}, data.m1);
  test.AddInput<float>("Moment_2", {3}, data.m2);

  // Verify AdamOptimizer outputs
  test.AddOutput<int64_t>("Update_Count_Out", {}, {4});
  test.AddOutput<float>("Moment_1_Out", {3}, data.m1_new);
  test.AddOutput<float>("Moment_2_Out", {3}, data.m2_new);
  test.AddOutput<float>("W_Out", {3}, data.w_new);

  test.AddAttribute("do_bias_correction", static_cast<int64_t>(0));
  test.AddAttribute("weight_decay_mode", static_cast<int64_t>(0));

  test.Run();
}

TEST(OptimizerTest, AdamOptimizerTest_Gradient) {
  OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
  AdamOptimizerInputOutput data;

  test.AddInput<float>("ETA", {}, data.eta);
  test.AddInput<int64_t>("Update_Count", {}, {3});
  test.AddInput<float>("W", {3}, data.w);
  test.AddInput<float>("G", {3}, data.g);
  test.AddInput<float>("Moment_1", {3}, data.m1);
  test.AddInput<float>("Moment_2", {3}, data.m2);

  // Verify AdamOptimizer outputs
  test.AddOutput<int64_t>("Update_Count_Out", {}, {4});
  test.AddOutput<float>("Moment_1_Out", {3}, data.m1_new);
  test.AddOutput<float>("Moment_2_Out", {3}, data.m2_new);
  test.AddOptionalOutputEdge<float>();
  test.AddOutput<float>("G_Out", {3}, data.g_new);

  test.AddAttribute("do_bias_correction", static_cast<int64_t>(0));
  test.AddAttribute("weight_decay_mode", static_cast<int64_t>(0));

  test.Run();
}

TEST(OptimizerTest, AdamBiasCorrection) {
  OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
  AdamOptimizerInputOutput data;

  test.AddInput<float>("ETA", {}, {1.f});
  test.AddInput<int64_t>("Update_Count", {}, {1});
  test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
  test.AddInput<float>("G", {3}, {0.4171f, 0.9485f, 1.2289f});
  test.AddInput<float>("Moment_1", {3}, {0.f, 0.f, 0.f});
  test.AddInput<float>("Moment_2", {3}, {0.f, 0.f, 0.f});

  test.AddOutput<int64_t>("Update_Count_Out", {}, {2});
  test.AddOutput<float>("Moment_1_Out", {3}, {0.0417f, 0.0949f, 0.1229f});
  test.AddOutput<float>("Moment_2_Out", {3}, {1.7400e-04f, 8.9966e-04f, 1.5102e-03f});
  test.AddOutput<float>("W_Out", {3}, {-1.4634f, -0.6416f, -1.2121f});

  test.AddAttribute("do_bias_correction", static_cast<int64_t>(1));
  test.AddAttribute("weight_decay_mode", static_cast<int64_t>(0));

  test.Run();
}

TEST(OptimizerTest, AdamWeightDecayMode0NoBiasCorrection) {
  OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
  AdamOptimizerInputOutput data;

  test.AddInput<float>("ETA", {}, {1.f});
  test.AddInput<int64_t>("Update_Count", {}, {1});
  test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
  test.AddInput<float>("G", {3}, {0.4171f, 0.9485f, 1.2289f});
  test.AddInput<float>("Moment_1", {3}, {0.f, 0.f, 0.f});
  test.AddInput<float>("Moment_2", {3}, {0.f, 0.f, 0.f});

  test.AddOutput<int64_t>("Update_Count_Out", {}, {2});
  test.AddOutput<float>("Moment_1_Out", {3}, {0.0417f, 0.0949f, 0.1229f});
  test.AddOutput<float>("Moment_2_Out", {3}, {1.7400e-04f, 8.9966e-04f, 1.5102e-03f});
  test.AddOutput<float>("W_Out", {3}, {-3.6210f, -2.8075f, -3.3723f});
  test.AddOutput<float>("G_Out", {3}, {-3.1576f, -3.1658f, -3.1601f});

  test.AddAttribute("do_bias_correction", static_cast<int64_t>(0));
  test.AddAttribute("lambda", 0.01f);
  test.AddAttribute("weight_decay_mode", static_cast<int64_t>(0));

  test.Run();
}

TEST(OptimizerTest, AdamWeightDecayMode0WithBiasCorrection) {
  OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
  AdamOptimizerInputOutput data;

  test.AddInput<float>("ETA", {}, {1.f});
  test.AddInput<int64_t>("Update_Count", {}, {1});
  test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
  test.AddInput<float>("G", {3}, {0.4171f, 0.9485f, 1.2289f});
  test.AddInput<float>("Moment_1", {3}, {0.f, 0.f, 0.f});
  test.AddInput<float>("Moment_2", {3}, {0.f, 0.f, 0.f});

  test.AddOutput<int64_t>("Update_Count_Out", {}, {2});
  test.AddOutput<float>("Moment_1_Out", {3}, {0.0417f, 0.0949f, 0.1229f});
  test.AddOutput<float>("Moment_2_Out", {3}, {1.7400e-04f, 8.9966e-04f, 1.5102e-03f});
  test.AddOutput<float>("W_Out", {3}, {-1.4587f, -0.6452f, -1.2099f});
  test.AddOutput<float>("G_Out", {3}, {-0.9954f, -1.0036f, -0.9979f});

  test.AddAttribute("do_bias_correction", static_cast<int64_t>(1));
  test.AddAttribute("lambda", 0.01f);
  test.AddAttribute("weight_decay_mode", static_cast<int64_t>(0));

  test.Run();
}

TEST(OptimizerTest, AdamWeightDecayMode1NoBiasCorrection) {
  OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
  AdamOptimizerInputOutput data;

  test.AddInput<float>("ETA", {}, {1.f});
  test.AddInput<int64_t>("Update_Count", {}, {1});
  test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
  test.AddInput<float>("G", {3}, {0.4171f, 0.9485f, 1.2289f});
  test.AddInput<float>("Moment_1", {3}, {0.f, 0.f, 0.f});
  test.AddInput<float>("Moment_2", {3}, {0.f, 0.f, 0.f});

  test.AddOutput<int64_t>("Update_Count_Out", {}, {2});
  test.AddOutput<float>("Moment_1_Out", {3}, {0.0417f, 0.0949f, 0.1229f});
  test.AddOutput<float>("Moment_2_Out", {3}, {1.7400e-04f, 8.9966e-04f, 1.5102e-03f});
  test.AddOutput<float>("W_Out", {3}, {-3.5894f, -2.7758f, -3.3406f});

  test.AddAttribute("do_bias_correction", static_cast<int64_t>(0));
  test.AddAttribute("lambda", 0.01f);
  test.AddAttribute("weight_decay_mode", static_cast<int64_t>(1));

  test.Run();
}

TEST(OptimizerTest, AdamWeightDecayMode1WithBiasCorrection) {
  OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
  AdamOptimizerInputOutput data;

  test.AddInput<float>("ETA", {}, {1.f});
  test.AddInput<int64_t>("Update_Count", {}, {1});
  test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
  test.AddInput<float>("G", {3}, {0.4171f, 0.9485f, 1.2289f});
  test.AddInput<float>("Moment_1", {3}, {0.f, 0.f, 0.f});
  test.AddInput<float>("Moment_2", {3}, {0.f, 0.f, 0.f});

  test.AddOutput<int64_t>("Update_Count_Out", {}, {2});
  test.AddOutput<float>("Moment_1_Out", {3}, {0.0417f, 0.0949f, 0.1229f});
  test.AddOutput<float>("Moment_2_Out", {3}, {1.7400e-04f, 8.9966e-04f, 1.5102e-03f});
  test.AddOutput<float>("W_Out", {3}, {-1.4488f, -0.6352f, -1.1999f});

  test.AddAttribute("do_bias_correction", static_cast<int64_t>(1));
  test.AddAttribute("lambda", 0.01f);
  test.AddAttribute("weight_decay_mode", static_cast<int64_t>(1));

  test.Run();
}

#if defined(USE_CUDA) || defined(USE_ROCM)

float GetGradientL2Norm(const std::vector<float>& gradient_vector) {
  float gradient_norm = 0.0f;
  for (const auto g_value : gradient_vector) {
    gradient_norm += g_value * g_value;
  }
  gradient_norm = std::sqrt(gradient_norm);
  return gradient_norm;
}

TEST(OptimizerTest, AdamOptimizerMixPrecisionTest) {
  OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
  AdamOptimizerInputOutput data;

  test.AddInput<MLFloat16>("ETA", {}, data.eta_half);
  test.AddInput<int64_t>("Update_Count", {}, {3});
  test.AddInput<float>("W", {3}, data.w);
  test.AddInput<MLFloat16>("G", {3}, data.g_half);
  test.AddInput<MLFloat16>("Moment_1", {3}, data.m1_half);
  test.AddInput<MLFloat16>("Moment_2", {3}, data.m2_half);

  // Verify AdamOptimizer outputs
  test.AddOutput<int64_t>("Update_Count_Out", {}, {4});
  test.AddOutput<MLFloat16>("Moment_1_Out", {3}, data.m1_new_half);
  test.AddOutput<MLFloat16>("Moment_2_Out", {3}, data.m2_new_half);
  test.AddOutput<float>("W_Out", {3}, data.w_new);

  test.AddAttribute("do_bias_correction", static_cast<int64_t>(0));
  test.AddAttribute("weight_decay_mode", static_cast<int64_t>(0));

  test.Run();
}

TEST(OptimizerTest, AdamOptimizerMixPrecision_FP16Weight_Test) {
  OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
  AdamOptimizerInputOutput data;

  test.AddInput<MLFloat16>("ETA", {}, data.eta_half);
  test.AddInput<int64_t>("Update_Count", {}, {3});
  test.AddInput<float>("W", {3}, data.w);
  test.AddInput<MLFloat16>("G", {3}, data.g_half);
  test.AddInput<MLFloat16>("Moment_1", {3}, data.m1_half);
  test.AddInput<MLFloat16>("Moment_2", {3}, data.m2_half);
  test.AddInput<MLFloat16>("FP16_W", {3}, data.w_half);

  // Verify AdamOptimizer outputs
  test.AddOutput<int64_t>("Update_Count_Out", {}, {4});
  test.AddOutput<MLFloat16>("Moment_1_Out", {3}, data.m1_new_half);
  test.AddOutput<MLFloat16>("Moment_2_Out", {3}, data.m2_new_half);
  test.AddOutput<float>("W_Out", {3}, data.w_new);
  test.AddOptionalOutputEdge<MLFloat16>();
  test.AddOutput<MLFloat16>("FP16_W_Out", {3}, data.w_new_half);

  test.AddAttribute("do_bias_correction", static_cast<int64_t>(0));
  test.AddAttribute("weight_decay_mode", static_cast<int64_t>(0));

  test.Run();
}

TEST(OptimizerTest, AdamOptimizerMixPrecision_FP16Weight_NoClipNorm_Test) {
  OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
  AdamOptimizerInputOutput data;

  test.AddInput<MLFloat16>("ETA", {}, data.eta_half);
  test.AddInput<int64_t>("Update_Count", {}, {3});
  test.AddInput<float>("W", {3}, data.w);
  test.AddInput<MLFloat16>("G", {3}, data.g_half);
  test.AddInput<MLFloat16>("Moment_1", {3}, data.m1_half);
  test.AddInput<MLFloat16>("Moment_2", {3}, data.m2_half);
  test.AddInput<MLFloat16>("FP16_W", {3}, data.w_half);
  test.AddInput<float>("loss_scale", {1}, {1.0f});
  // grad clipping should not take effect because default max_norm is 1.0f
  test.AddInput<float>("grad_norm", {1}, {0.01f});
  // Verify AdamOptimizer outputs
  test.AddOutput<int64_t>("Update_Count_Out", {}, {4});
  test.AddOutput<MLFloat16>("Moment_1_Out", {3}, data.m1_new_half);
  test.AddOutput<MLFloat16>("Moment_2_Out", {3}, data.m2_new_half);
  test.AddOutput<float>("W_Out", {3}, data.w_new);
  test.AddOptionalOutputEdge<MLFloat16>();
  test.AddOutput<MLFloat16>("FP16_W_Out", {3}, data.w_new_half);

  test.AddAttribute("do_bias_correction", static_cast<int64_t>(0));
  test.AddAttribute("weight_decay_mode", static_cast<int64_t>(0));
  test.Run();
}

TEST(OptimizerTest, AdamOptimizerMixPrecision_FP16Weight_ClipNorm_Test) {
  OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
  AdamOptimizerInputOutput data;

  // Expected FP32 Outputs
  std::vector<float> m1_new = {0.13f, 0.23f, 0.33f};
  std::vector<float> m2_new = {0.3997f, 0.4998f, 0.6001f};
  std::vector<float> w_new = {0.8972168f, 1.8369141f, 2.7871094f};
  // FP16 Outputs
  std::vector<MLFloat16> m1_new_half;
  std::vector<MLFloat16> m2_new_half;
  std::vector<MLFloat16> w_new_half;

  m1_new_half.resize(m1_new.size());
  m2_new_half.resize(m2_new.size());
  w_new_half.resize(w_new.size());
  ConvertFloatToMLFloat16(m1_new.data(), m1_new_half.data(), int(m1_new.size()));
  ConvertFloatToMLFloat16(m2_new.data(), m2_new_half.data(), int(m2_new.size()));
  ConvertFloatToMLFloat16(w_new.data(), w_new_half.data(), int(w_new.size()));

  test.AddInput<MLFloat16>("ETA", {}, data.eta_half);
  test.AddInput<int64_t>("Update_Count", {}, {3});
  test.AddInput<float>("W", {3}, data.w);
  test.AddInput<MLFloat16>("G", {3}, data.g_half);
  test.AddInput<MLFloat16>("Moment_1", {3}, data.m1_half);
  test.AddInput<MLFloat16>("Moment_2", {3}, data.m2_half);
  test.AddInput<MLFloat16>("FP16_W", {3}, data.w_half);
  test.AddInput<float>("loss_scale", {1}, {1.0f});
  test.AddInput<float>("grad_norm", {1}, {0.01f});
  // Verify AdamOptimizer outputs
  test.AddOutput<int64_t>("Update_Count_Out", {}, {4});
  test.AddOutput<MLFloat16>("Moment_1_Out", {3}, m1_new_half);
  test.AddOutput<MLFloat16>("Moment_2_Out", {3}, m2_new_half);
  test.AddOutput<float>("W_Out", {3}, w_new);
  test.AddOptionalOutputEdge<MLFloat16>();
  test.AddOutput<MLFloat16>("FP16_W_Out", {3}, w_new_half);

  test.AddAttribute("do_bias_correction", static_cast<int64_t>(0));
  test.AddAttribute("weight_decay_mode", static_cast<int64_t>(0));
  test.AddAttribute("max_norm_clip", 0.001f);
  test.Run();
}

TEST(OptimizerTest, AdamOptimizerMixPrecision_FP16Weight_SkipUpdate_Test) {
  OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
  AdamOptimizerInputOutput data;

  test.AddInput<MLFloat16>("ETA", {}, data.eta_half);
  test.AddInput<int64_t>("Update_Count", {}, {3});
  test.AddInput<float>("W", {3}, data.w);
  test.AddInput<MLFloat16>("G", {3}, data.g_half);
  test.AddInput<MLFloat16>("Moment_1", {3}, data.m1_half);
  test.AddInput<MLFloat16>("Moment_2", {3}, data.m2_half);
  test.AddInput<MLFloat16>("FP16_W", {3}, data.w_half);
  test.AddInput<float>("loss_scale", {1}, {1.0f});
  // grad clipping should not take effect
  test.AddInput<float>("grad_norm", {1}, {0.01f});
  test.AddInput<bool>("DoUpdate", {1}, {false});

  // Verify AdamOptimizer outputs
  test.AddOutput<int64_t>("Update_Count_Out", {}, {3});
  test.AddOutput<MLFloat16>("Moment_1_Out", {3}, data.m1_half);
  test.AddOutput<MLFloat16>("Moment_2_Out", {3}, data.m2_half);
  test.AddOutput<float>("W_Out", {3}, data.w);
  test.AddOptionalOutputEdge<MLFloat16>();
  test.AddOutput<MLFloat16>("FP16_W_Out", {3}, data.w_half);

  test.AddAttribute("do_bias_correction", static_cast<int64_t>(0));
  test.AddAttribute("weight_decay_mode", static_cast<int64_t>(0));

  test.Run();
}

TEST(OptimizerTest, AdamOptimizerMixPrecisionTestFloatEta) {
  OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
  AdamOptimizerInputOutput data;

  test.AddInput<float>("ETA", {}, data.eta);
  test.AddInput<int64_t>("Update_Count", {}, {3});
  test.AddInput<float>("W", {3}, data.w);
  test.AddInput<MLFloat16>("G", {3}, data.g_half);
  test.AddInput<MLFloat16>("Moment_1", {3}, data.m1_half);
  test.AddInput<MLFloat16>("Moment_2", {3}, data.m2_half);

  // Verify AdamOptimizer outputs
  test.AddOutput<int64_t>("Update_Count_Out", {}, {4});
  test.AddOutput<MLFloat16>("Moment_1_Out", {3}, data.m1_new_half);
  test.AddOutput<MLFloat16>("Moment_2_Out", {3}, data.m2_new_half);
  test.AddOutput<float>("W_Out", {3}, data.w_new);

  test.AddAttribute("do_bias_correction", static_cast<int64_t>(0));
  test.AddAttribute("weight_decay_mode", static_cast<int64_t>(0));

  test.Run();
}

TEST(OptimizerTest, AdamOptimizerMixPrecisionTest_Gradient) {
  OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
  AdamOptimizerInputOutput data;

  test.AddInput<float>("ETA", {}, data.eta);
  test.AddInput<int64_t>("Update_Count", {}, {3});
  test.AddInput<float>("W", {3}, data.w);
  test.AddInput<MLFloat16>("G", {3}, data.g_half);
  test.AddInput<MLFloat16>("Moment_1", {3}, data.m1_half);
  test.AddInput<MLFloat16>("Moment_2", {3}, data.m2_half);

  // Verify AdamOptimizer outputs
  test.AddOutput<int64_t>("Update_Count_Out", {}, {4});
  test.AddOutput<MLFloat16>("Moment_1_Out", {3}, data.m1_new_half);
  test.AddOutput<MLFloat16>("Moment_2_Out", {3}, data.m2_new_half);
  test.AddOptionalOutputEdge<float>();
  test.AddOutput<MLFloat16>("G_Out", {3}, data.g_new_half);

  test.AddAttribute("do_bias_correction", static_cast<int64_t>(0));

  test.Run();
}

// This helper function is a CPU-based LAMB optimizer
// implementation. It mainly focuses on readability.
void compute_lamb(
    const std::vector<int64_t> shape,
    /* weights */ const std::vector<float>& w,
    /* gradient */ const std::vector<float>& g,
    /* momentum */ const std::vector<float>& m,
    /* 2nd-order momentum */ const std::vector<float>& v,
    const float eta,
    const float lambda,
    const float alpha,
    const float beta,
    const float epsilon,
    const float max_norm_clip,
    /* updated weights */ std::vector<float>& w_new,
    /* updated gradients */ std::vector<float>& g_new,
    /* updated momentum */ std::vector<float>& m_new,
    /* updated 2nd-order momentum */ std::vector<float>& v_new,
    const int64_t step = 0,
    const float loss_scale = 1.0f,
    const float* p_scaled_g_norm = nullptr,
    const float ratio_min = -std::numeric_limits<float>::infinity(),
    const float ratio_max = std::numeric_limits<float>::infinity()) {
  // Element counts of all vector-typed arguments.
  const int64_t size = std::accumulate(
      shape.begin(),
      shape.end(),
      (int64_t)1,
      std::multiplies<int64_t>());

  // Buffer to store update direction.
  std::vector<float> r(size, 0.0f);

  float scaled_g_scaling_factor = loss_scale;
  if (p_scaled_g_norm != nullptr) {
    const float scaled_g_max_norm = loss_scale * max_norm_clip;
    if (*p_scaled_g_norm > scaled_g_max_norm) {
      scaled_g_scaling_factor = *p_scaled_g_norm / max_norm_clip;
    }
  }

  const float alpha_correction = step > 0 ? 1.f - std::pow(alpha, static_cast<float>(step)) : 1.f;
  const float beta_correction = step > 0 ? 1.f - std::pow(beta, static_cast<float>(step)) : 1.f;

  // Compute new 1st-, 2nd-order momentums, and the update direction.
  for (int i = 0; i < size; ++i) {
    const float g_scaled = g[i] / scaled_g_scaling_factor;
    m_new[i] = alpha * m[i] + (1.0f - alpha) * g_scaled;
    v_new[i] = beta * v[i] + (1.0f - beta) * g_scaled * g_scaled;
    const float m_new_tmp = m_new[i] / alpha_correction;
    const float v_new_tmp = v_new[i] / beta_correction;
    r[i] = lambda * w[i] + m_new_tmp / (std::sqrt(v_new_tmp) + epsilon);
  }

  // Compute squared sum of all elements. Note that Eigen sqrt could lead to significant
  // numerical error so we use std::sqrt. The std::inner_product produces wrong result
  // when std::inner_product(r.begin(), r.end(), r.begin(), 0) so we just use a loop below.
  float r_norm = 0.0f;
  float w_norm = 0.0f;
  for (int i = 0; i < size; ++i) {
    r_norm += r[i] * r[i];
    w_norm += w[i] * w[i];
  }

  r_norm = std::sqrt(r_norm);
  w_norm = std::sqrt(w_norm);

  float ratio = (w_norm != 0.0f && r_norm != 0.0f) ? w_norm / r_norm : 1.0f;

  if (ratio > ratio_max) {
    ratio = ratio_max;
  }

  if (ratio < ratio_min) {
    ratio = ratio_min;
  }

  ratio *= eta;

  // Compute the new weight.
  for (int64_t i = 0; i < size; ++i) {
    g_new[i] = -ratio * r[i];
    w_new[i] = w[i] + g_new[i];
  }
}

template <typename T1, typename T2, typename T3, typename T4>
void run_lamb_test_with_baseline(
    const std::vector<int64_t>& shape,
    const std::vector<T1>& eta,
    const std::vector<T2>& w,
    const std::vector<T3>& g,
    const std::vector<T4>& m,
    const std::vector<T4>& v,
    const float alpha,
    const float beta,
    const float lambda,
    const float epsilon,
    const float max_norm,
    const std::vector<T2>& w_new,
    const std::vector<T3>& g_new,
    const std::vector<T4>& m_new,
    const std::vector<T4>& v_new,
    const std::vector<MLFloat16>& w_half = {},
    const std::vector<MLFloat16>& w_new_half = {},
    const bool do_update = true,
    const int64_t step = 0,
    const float loss_scale = 1.0f,
    const float* p_g_norm = nullptr,
    const float ratio_min = -std::numeric_limits<float>::infinity(),
    const float ratio_max = std::numeric_limits<float>::infinity()) {
  OpTester test("LambOptimizer", 1, onnxruntime::kMSDomain, true);

  test.AddInput<bool>("update_signal", {1}, {do_update});
  test.AddInput<T2>("loss_scale", {}, {loss_scale});
  if (p_g_norm == nullptr) {
    test.AddOptionalInputEdge<T2>();
  } else {
    test.AddInput<T2>("gradient_norm", {}, {T2(*p_g_norm)});
  }
  test.AddInput<T1>("ETA", {1}, eta);
  if (step > 0) {
    test.AddInput<int64_t>("Step", {}, {step});
  } else {
    test.AddOptionalInputEdge<int64_t>();
  }
  test.AddInput<T2>("W", shape, w);
  test.AddInput<T3>("G", shape, g);
  test.AddInput<T4>("Moment_1", shape, m);
  test.AddInput<T4>("Moment_2", shape, v);
  if (!w_half.empty()) {
    test.AddInput<MLFloat16>("FP16_W", shape, w_half);
  } else {
    test.AddOptionalInputEdge<MLFloat16>();
  }

  test.AddAttribute("alpha", std::vector<float>(1, alpha));
  test.AddAttribute("beta", std::vector<float>(1, beta));
  test.AddAttribute("lambda", std::vector<float>(1, lambda));
  test.AddAttribute("epsilon", std::vector<float>(1, epsilon));
  test.AddAttribute("max_norm_clip", std::vector<float>(1, max_norm));
  test.AddAttribute("ratio_min", ratio_min);
  test.AddAttribute("ratio_max", ratio_max);

  if (step > 0) {
    test.AddOutput<int64_t>("Step_Out", {}, {do_update ? step + 1 : step});
  } else {
    test.AddOptionalOutputEdge<int64_t>();
  }
  if (!w_new.empty()) {
    test.AddOutput<T2>("W_Out", shape, w_new);
  } else {
    test.AddOptionalOutputEdge<T2>();
  }
  if (!g_new.empty()) {
    test.AddOutput<T3>("G_Out", shape, g_new);
  } else {
    test.AddOptionalOutputEdge<T3>();
  }
  test.AddOutput<T4>("Moment_1_Out", shape, m_new);
  test.AddOutput<T4>("Moment_2_Out", shape, v_new);
  if (!w_new_half.empty()) {
    test.AddOutput<MLFloat16>("FP16_W_Out", shape, w_new_half);
  } else {
    test.AddOptionalOutputEdge<MLFloat16>();
  }

  test.Run();
}

template <typename T1, typename T2, typename T3, typename T4>
void run_multi_tensor_lamb_test_with_baseline(
    const std::vector<std::vector<int64_t>>& shapes,
    const T1 eta,
    const std::vector<std::vector<T2>>& ws,
    const std::vector<std::vector<T3>>& gs,
    const std::vector<std::vector<T4>>& ms,
    const std::vector<std::vector<T4>>& vs,
    const std::vector<float>& alphas,
    const std::vector<float>& betas,
    const std::vector<float>& lambdas,
    const std::vector<float>& epsilons,
    const std::vector<float>& max_norms,
    const std::vector<std::vector<T2>>& w_news,
    const std::vector<std::vector<T3>>& g_news,
    const std::vector<std::vector<T4>>& m_news,
    const std::vector<std::vector<T4>>& v_news,
    const std::vector<std::vector<MLFloat16>>& w_halfs = {},
    const std::vector<std::vector<MLFloat16>>& w_new_halfs = {},
    const bool do_update = true,
    const int64_t step = 0,
    const float loss_scale = 1.0f,
    const float* p_g_norm = nullptr,
    const float ratio_min = -std::numeric_limits<float>::infinity(),
    const float ratio_max = std::numeric_limits<float>::infinity()) {
  OpTester test("LambOptimizer", 1, onnxruntime::kMSDomain, true);

  ORT_ENFORCE(shapes.size() == ws.size());
  ORT_ENFORCE(shapes.size() == gs.size());
  ORT_ENFORCE(shapes.size() == ms.size());
  ORT_ENFORCE(shapes.size() == vs.size());
  ORT_ENFORCE(shapes.size() == alphas.size());
  ORT_ENFORCE(shapes.size() == betas.size());
  ORT_ENFORCE(shapes.size() == lambdas.size());
  ORT_ENFORCE(shapes.size() == epsilons.size());
  ORT_ENFORCE(shapes.size() == max_norms.size());
  if (!w_news.empty()) {
    ORT_ENFORCE(shapes.size() == w_news.size());
  }
  if (!g_news.empty()) {
    ORT_ENFORCE(shapes.size() == g_news.size());
  }
  ORT_ENFORCE(shapes.size() == m_news.size());
  ORT_ENFORCE(shapes.size() == v_news.size());
  if (!w_halfs.empty()) {
    ORT_ENFORCE(shapes.size() == w_halfs.size());
  }
  if (!w_new_halfs.empty()) {
    ORT_ENFORCE(shapes.size() == w_new_halfs.size());
  }

  const int group_count = static_cast<int>(ws.size());

  test.AddInput<bool>("update_signal", {}, {do_update});
  test.AddInput<T2>("loss_scale", {}, {loss_scale});
  if (p_g_norm == nullptr) {
    test.AddOptionalInputEdge<T2>();
  } else {
    test.AddInput<float>("gradient_norm", {}, {T2(*p_g_norm)});
  }
  test.AddInput<T1>("ETA", {}, {eta});
  if (step > 0) {
    test.AddInput<int64_t>("Step", {}, {step});
    test.AddOutput<int64_t>("Step_Out", {}, {do_update ? step + 1 : step});
  } else {
    test.AddOptionalInputEdge<int64_t>();
    test.AddOptionalOutputEdge<int64_t>();
  }
  for (int i = 0; i < group_count; ++i) {
    std::string w_name = "W_" + std::to_string(i);
    std::string g_name = "G_" + std::to_string(i);
    std::string m1_name = "Moment_1_" + std::to_string(i);
    std::string m2_name = "Moment_2_" + std::to_string(i);
    std::string w_fp16_name = "FP16_W_" + std::to_string(i);
    std::string w_new_name = "W_Out_" + std::to_string(i);
    std::string g_new_name = "G_Out_" + std::to_string(i);
    std::string m1_new_name = "Moment_1_Out_" + std::to_string(i);
    std::string m2_new_name = "Moment_2_Out_" + std::to_string(i);
    std::string w_fp16_new_name = "FP16_W_Out_" + std::to_string(i);

    test.AddInput<T2>(w_name.c_str(), shapes[i], ws[i]);
    test.AddInput<T3>(g_name.c_str(), shapes[i], gs[i]);
    test.AddInput<T4>(m1_name.c_str(), shapes[i], ms[i]);
    test.AddInput<T4>(m2_name.c_str(), shapes[i], vs[i]);
    if (!w_halfs.empty() && !w_halfs[i].empty()) {
      test.AddInput<MLFloat16>(w_fp16_name.c_str(), shapes[i], w_halfs[i]);
    } else {
      test.AddOptionalInputEdge<MLFloat16>();
    }

    if (!w_news.empty() && !w_news[i].empty()) {
      test.AddOutput<T2>(w_new_name.c_str(), shapes[i], w_news[i]);
    } else {
      test.AddOptionalOutputEdge<T2>();
    }
    if (!g_news.empty() && !g_news[i].empty()) {
      test.AddOutput<T3>(g_new_name.c_str(), shapes[i], g_news[i]);
    } else {
      test.AddOptionalOutputEdge<T3>();
    }
    test.AddOutput<T4>(m1_new_name.c_str(), shapes[i], m_news[i]);
    test.AddOutput<T4>(m2_new_name.c_str(), shapes[i], v_news[i]);
    if (!w_new_halfs.empty() && !w_new_halfs[i].empty()) {
      test.AddOutput<MLFloat16>(w_fp16_new_name.c_str(), shapes[i], w_new_halfs[i]);
    } else {
      test.AddOptionalOutputEdge<MLFloat16>();
    }
  }

  test.AddAttribute("alpha", alphas);
  test.AddAttribute("beta", betas);
  test.AddAttribute("lambda", lambdas);
  test.AddAttribute("epsilon", epsilons);
  test.AddAttribute("max_norm_clip", max_norms);
  test.AddAttribute("ratio_min", ratio_min);
  test.AddAttribute("ratio_max", ratio_max);

  test.Run();
}

// Lamb test without baseline. This function computes
// baseline via an internal function and then invoke
// run_lamb_test_with_baseline(...) to check the result.
void run_multi_tensor_lamb_test(
    const std::vector<std::vector<int64_t>> shapes,
    const float eta,
    const std::vector<std::vector<float>> ws,
    const std::vector<std::vector<float>> gs,
    const std::vector<std::vector<float>> ms,
    const std::vector<std::vector<float>> vs,
    const std::vector<float> lambdas,
    const std::vector<float> alphas,
    const std::vector<float> betas,
    const std::vector<float> epsilons,
    const std::vector<float> max_norms,
    const int64_t step = 0,
    const float loss_scale = 1.0f,
    const float* p_scaled_g_norm = nullptr,
    const float ratio_min = -std::numeric_limits<float>::infinity(),
    const float ratio_max = std::numeric_limits<float>::infinity()) {
  // Check if parallel vectors have the same length.
  ORT_ENFORCE(shapes.size() == ws.size());
  ORT_ENFORCE(shapes.size() == gs.size());
  ORT_ENFORCE(shapes.size() == ms.size());
  ORT_ENFORCE(shapes.size() == vs.size());
  ORT_ENFORCE(shapes.size() == alphas.size());
  ORT_ENFORCE(shapes.size() == betas.size());
  ORT_ENFORCE(shapes.size() == lambdas.size());
  ORT_ENFORCE(shapes.size() == epsilons.size());
  ORT_ENFORCE(shapes.size() == max_norms.size());

  const int group_count = static_cast<int>(ws.size());

  // Output buffers of the optimizer.
  std::vector<std::vector<float>> w_news(group_count);
  std::vector<std::vector<float>> g_news(group_count);
  std::vector<std::vector<float>> m_news(group_count);
  std::vector<std::vector<float>> v_news(group_count);

  for (int i = 0; i < group_count; ++i) {
    w_news[i] = std::vector<float>(ws[i].size(), 0.f);
    g_news[i] = std::vector<float>(gs[i].size(), 0.f);
    m_news[i] = std::vector<float>(ms[i].size(), 0.f);
    v_news[i] = std::vector<float>(vs[i].size(), 0.f);

    // Invoke LAMB's reference implementation to compute baseline output.
    compute_lamb(
        shapes[i], ws[i], gs[i], ms[i], vs[i],
        eta, lambdas[i], alphas[i], betas[i], epsilons[i], max_norms[i],
        w_news[i], g_news[i], m_news[i], v_news[i], step, loss_scale, p_scaled_g_norm,
        ratio_min, ratio_max);
  }

  // Create tests to make sure the output is correct.

  // Output new weights.
  run_multi_tensor_lamb_test_with_baseline(
      shapes, eta,
      ws, gs, ms, vs,
      alphas, betas, lambdas, epsilons, max_norms,
      w_news, {}, m_news, v_news, {}, {}, true, step, loss_scale, p_scaled_g_norm,
      ratio_min, ratio_max);

  // Output new gradients.
  run_multi_tensor_lamb_test_with_baseline(
      shapes, eta,
      ws, gs, ms, vs,
      alphas, betas, lambdas, epsilons, max_norms,
      {}, g_news, m_news, v_news, {}, {}, true, step, loss_scale, p_scaled_g_norm,
      ratio_min, ratio_max);
}

void run_lamb_mix_precision_test(
    const std::vector<int64_t>& shape,
    const std::vector<float>& eta,
    const std::vector<float>& w,
    const std::vector<float>& g,
    const std::vector<float>& m,
    const std::vector<float>& v,
    const float lambda,
    const float alpha,
    const float beta,
    const float epsilon,
    const float max_norm,
    const int64_t step = 0,
    const float loss_scale = 1.0f,
    const float* p_g_norm = nullptr) {
  std::vector<float> w_new(w.size(), 0);
  std::vector<float> g_new(g.size(), 0);
  std::vector<float> m_new(m.size(), 0);
  std::vector<float> v_new(v.size(), 0);

  // Invoke LAMB's reference implementation to compute output.
  compute_lamb(
      shape, w, g, m, v,
      eta[0], lambda, alpha, beta, epsilon, max_norm,
      w_new, g_new, m_new, v_new, step, loss_scale, p_g_norm);

  std::vector<MLFloat16> eta_half(eta.size());
  std::vector<MLFloat16> g_half(w.size());
  std::vector<MLFloat16> m_half(w.size());
  std::vector<MLFloat16> v_half(w.size());
  std::vector<MLFloat16> w_half(w.size());
  ConvertFloatToMLFloat16(eta.data(), eta_half.data(), int(eta.size()));
  ConvertFloatToMLFloat16(g.data(), g_half.data(), int(g.size()));
  ConvertFloatToMLFloat16(m.data(), m_half.data(), int(m.size()));
  ConvertFloatToMLFloat16(v.data(), v_half.data(), int(v.size()));
  ConvertFloatToMLFloat16(w.data(), w_half.data(), int(w.size()));

  std::vector<MLFloat16> m_new_half(m_new.size());
  std::vector<MLFloat16> v_new_half(v_new.size());
  std::vector<MLFloat16> w_new_half(w_new.size());
  std::vector<MLFloat16> g_new_half(g_new.size());
  ConvertFloatToMLFloat16(m_new.data(), m_new_half.data(), int(m_new.size()));
  ConvertFloatToMLFloat16(v_new.data(), v_new_half.data(), int(v_new.size()));
  ConvertFloatToMLFloat16(w_new.data(), w_new_half.data(), int(w_new.size()));
  ConvertFloatToMLFloat16(g_new.data(), g_new_half.data(), int(g_new.size()));

  // Half momentums, without fp16 weight
  run_lamb_test_with_baseline(
      shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, max_norm,
      w_new, {}, m_new_half, v_new_half, {}, {}, true, step, loss_scale, p_g_norm);

  // Float momentums, without fp16 weight
  run_lamb_test_with_baseline(
      shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
      w_new, {}, m_new, v_new, {}, {}, true, step, loss_scale, p_g_norm);

  // Half momentums, with fp16 weight
  run_lamb_test_with_baseline(
      shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, max_norm,
      w_new, {}, m_new_half, v_new_half, {}, {}, true, step, loss_scale, p_g_norm);

  // Float momentums, with fp16 weight
  run_lamb_test_with_baseline(
      shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
      w_new, {}, m_new, v_new, w_half, w_new_half, true, step, loss_scale, p_g_norm);

  // Half momentums, with fp16 weight, skip weight update
  run_lamb_test_with_baseline(
      shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, max_norm,
      w, {}, m_half, v_half, w_half, w_half, false, step, loss_scale, p_g_norm);

  // Float momentums, with fp16 weight, skip weight update
  run_lamb_test_with_baseline(
      shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
      w, {}, m, v, w_half, w_half, false, step, loss_scale, p_g_norm);

  // Float eta, float momentums, with fp16 weight
  run_lamb_test_with_baseline(
      shape, eta, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
      w_new, {}, m_new, v_new, w_half, w_new_half, true, step, loss_scale, p_g_norm);

  // Float eta, float momentums, with fp16 weight, skip weight update
  run_lamb_test_with_baseline(
      shape, eta, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
      w, {}, m, v, w_half, w_half, false, step, loss_scale, p_g_norm);

  // Float momentums, without fp16 weight, output gradients only
  run_lamb_test_with_baseline(
      shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
      {}, g_new_half, m_new, v_new, {}, {}, true, step, loss_scale, p_g_norm);

  // Float momentums, with fp16 weight, output gradients only
  run_lamb_test_with_baseline(
      shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
      {}, g_new_half, m_new, v_new, w_half, {}, true, step, loss_scale, p_g_norm);

  // Float momentums, with fp16 weight, output gradients only, skip weight update
  run_lamb_test_with_baseline(
      shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
      {}, g_half, m, v, w_half, {}, false, step, loss_scale, p_g_norm);
}

// A optimizer test with an 2-element vector.
TEST(OptimizerTest, LambOptimizerTestVector) {
  // Input tensors and attributes.
  const std::vector<int64_t> shape = {2};
  constexpr float eta = 0.5f;
  const std::vector<float> w = {1.0f, 2.0f};
  const std::vector<float> g = {3.0f, 4.0f};
  const std::vector<float> m = {-1.0f, -2.0f};
  const std::vector<float> v = {2.0f, 1.0f};
  constexpr float lambda = 0.5f;
  constexpr float alpha = 0.2f;
  constexpr float beta = 0.8f;
  constexpr float epsilon = 1e-6f;
  constexpr float max_norm = 1.0f;

  constexpr int64_t step = 0;
  constexpr float loss_scale = 1.f;
  constexpr float scaled_g_norm = 1.f;
  run_multi_tensor_lamb_test(
      {shape},
      eta,
      {w},
      {g},
      {m},
      {v},
      {lambda},
      {alpha},
      {beta},
      {epsilon},
      {max_norm},
      step,
      loss_scale,
      &scaled_g_norm);
}

TEST(OptimizerTest, LambOptimizerTestVectorWithZeroWeight) {
  // Input tensors and attributes.
  const std::vector<int64_t> shape = {2};
  constexpr float eta = 0.5f;
  const std::vector<float> w = {0.0f, 0.0f};
  const std::vector<float> g = {1.0f, -1.0f};
  const std::vector<float> m = {-1.0f, -2.0f};
  const std::vector<float> v = {2.0f, 1.0f};
  constexpr float lambda = 0.5f;
  constexpr float alpha = 0.2f;
  constexpr float beta = 0.8f;
  constexpr float epsilon = 1e-6f;
  constexpr float max_norm = 1.0f;
  constexpr int64_t step = 0;
  constexpr float loss_scale = 1.f;
  constexpr float scaled_g_norm = 1.f;
  run_multi_tensor_lamb_test(
      {shape},
      eta,
      {w},
      {g},
      {m},
      {v},
      {lambda},
      {alpha},
      {beta},
      {epsilon},
      {max_norm},
      step,
      loss_scale,
      &scaled_g_norm);
}

TEST(OptimizerTest, LambOptimizerRatioMin) {
  // Input tensors and attributes.
  const std::vector<int64_t> shape = {2};
  constexpr float eta = 0.5f;
  const std::vector<float> w = {-1.0f, 1.0f};
  const std::vector<float> g = {1.0f, -1.0f};
  const std::vector<float> m = {-1.0f, -2.0f};
  const std::vector<float> v = {2.0f, 1.0f};
  constexpr float lambda = 0.5f;
  constexpr float alpha = 0.2f;
  constexpr float beta = 0.8f;
  constexpr float epsilon = 1e-6f;
  constexpr float max_norm = 1.0f;
  constexpr int64_t step = 0;
  constexpr float loss_scale = 1.f;
  constexpr float scaled_g_norm = 1.f;
  constexpr float ratio_min = -std::numeric_limits<float>::infinity();
  constexpr float ratio_max = 0.1f;

  run_multi_tensor_lamb_test(
      {shape},
      eta,
      {w},
      {g},
      {m},
      {v},
      {lambda},
      {alpha},
      {beta},
      {epsilon},
      {max_norm},
      step,
      loss_scale,
      &scaled_g_norm,
      ratio_min,
      ratio_max);
}

TEST(OptimizerTest, LambOptimizerRatioMax) {
  // Input tensors and attributes.
  const std::vector<int64_t> shape = {2};
  constexpr float eta = 0.5f;
  const std::vector<float> w = {0.0001f, -0.0001f};
  const std::vector<float> g = {1.0f, -1.0f};
  const std::vector<float> m = {-1.0f, -2.0f};
  const std::vector<float> v = {2.0f, 1.0f};
  constexpr float lambda = 0.5f;
  constexpr float alpha = 0.2f;
  constexpr float beta = 0.8f;
  constexpr float epsilon = 1e-6f;
  constexpr float max_norm = 1.0f;
  constexpr int64_t step = 0;
  constexpr float loss_scale = 1.f;
  constexpr float scaled_g_norm = 1.f;
  constexpr float ratio_min = 1.0f;
  constexpr float ratio_max = std::numeric_limits<float>::infinity();

  run_multi_tensor_lamb_test(
      {shape},
      eta,
      {w},
      {g},
      {m},
      {v},
      {lambda},
      {alpha},
      {beta},
      {epsilon},
      {max_norm},
      step,
      loss_scale,
      &scaled_g_norm,
      ratio_min,
      ratio_max);
}

TEST(OptimizerTest, LambOptimizerTestBiasCorrectionFirst) {
  // Input tensors and attributes.
  const std::vector<int64_t> shape = {2};
  constexpr float eta = 0.5f;
  const std::vector<float> w = {1.0f, 2.0f};
  const std::vector<float> g = {3.0f, 4.0f};
  const std::vector<float> m = {-1.0f, -2.0f};
  const std::vector<float> v = {2.0f, 1.0f};
  constexpr float lambda = 0.5f;
  constexpr float alpha = 0.2f;
  constexpr float beta = 0.8f;
  constexpr float epsilon = 1e-6f;
  constexpr float max_norm = 1.0f;

  constexpr int64_t step = 1;
  constexpr float loss_scale = 1.f;
  constexpr float scaled_g_norm = 1.f;
  run_multi_tensor_lamb_test(
      {shape},
      eta,
      {w},
      {g},
      {m},
      {v},
      {lambda},
      {alpha},
      {beta},
      {epsilon},
      {max_norm},
      step,
      loss_scale,
      &scaled_g_norm);
}

TEST(OptimizerTest, LambOptimizerTestBiasCorrectionThird) {
  // Input tensors and attributes.
  const std::vector<int64_t> shape = {2};
  constexpr float eta = 0.5f;
  const std::vector<float> w = {1.0f, 2.0f};
  const std::vector<float> g = {3.0f, 4.0f};
  const std::vector<float> m = {-1.0f, -2.0f};
  const std::vector<float> v = {2.0f, 1.0f};
  constexpr float lambda = 0.5f;
  constexpr float alpha = 0.2f;
  constexpr float beta = 0.8f;
  constexpr float epsilon = 1e-6f;
  constexpr float max_norm = 1.0f;

  constexpr int64_t step = 3;
  constexpr float loss_scale = 1.f;
  constexpr float scaled_g_norm = 1.f;
  run_multi_tensor_lamb_test(
      {shape},
      eta,
      {w},
      {g},
      {m},
      {v},
      {lambda},
      {alpha},
      {beta},
      {epsilon},
      {max_norm},
      step,
      loss_scale,
      &scaled_g_norm);
}

// A optimizer test with an 2-by-1-by-1-by-1 tensor.
TEST(OptimizerTest, LambOptimizerTest4DTensor) {
  // Input tensors and attributes.
  const std::vector<int64_t> shape = {2, 1, 1, 1};
  constexpr float eta = 0.5f;
  const std::vector<float> w = {1.0f, 2.0f};
  const std::vector<float> g = {3.0f, 4.0f};
  const std::vector<float> m = {-1.0f, -2.0f};
  const std::vector<float> v = {2.0f, 1.0f};
  constexpr float lambda = 0.5f;
  constexpr float alpha = 0.2f;
  constexpr float beta = 0.8f;
  constexpr float epsilon = 1e-6f;
  constexpr float max_norm = 1.0f;

  constexpr int64_t step = 0;
  constexpr float loss_scale = 1.f;
  constexpr float scaled_g_norm = 1.f;
  run_multi_tensor_lamb_test(
      {shape},
      eta,
      {w},
      {g},
      {m},
      {v},
      {lambda},
      {alpha},
      {beta},
      {epsilon},
      {max_norm},
      step,
      loss_scale,
      &scaled_g_norm);
}

// A optimizer test with an 2-by-3 tensor.
TEST(OptimizerTest, LambOptimizerTest2by3Tensor) {
  // Input tensors and attributes.
  const std::vector<int64_t> shape = {2, 3};
  constexpr float eta = 0.5f;
  const std::vector<float> w = {1.0f, 2.0f, 1.0f, 1.0f, 2.0f, 2.0f};
  const std::vector<float> g = {3.0f, 4.0f, 3.0f, 3.0f, 4.0f, 4.0f};
  const std::vector<float> m = {-1.0f, -2.0f, 2.0f, 1.0f, 1.0f, -2.0f};
  const std::vector<float> v = {1.0f, 1.0f, 5.0f, 5.0f, 6.0f, 6.0f};
  constexpr float lambda = 0.5f;
  constexpr float alpha = 0.2f;
  constexpr float beta = 0.8f;
  constexpr float epsilon = 1e-6f;
  constexpr float max_norm = 1.0f;

  constexpr int64_t step = 0;
  constexpr float loss_scale = 1.f;
  constexpr float scaled_g_norm = 1.f;
  run_multi_tensor_lamb_test(
      {shape},
      eta,
      {w},
      {g},
      {m},
      {v},
      {lambda},
      {alpha},
      {beta},
      {epsilon},
      {max_norm},
      step,
      loss_scale,
      &scaled_g_norm);
}

// A optimizer test with an 1-element tensor.
TEST(OptimizerTest, LambOptimizerTestScalar) {
  // Input tensors and attributes.
  const std::vector<int64_t> shape = {(int64_t)1};
  constexpr float eta = 0.5f;
  const std::vector<float> w = {1.0f};
  const std::vector<float> g = {3.0f};
  const std::vector<float> m = {-10.0f};
  const std::vector<float> v = {1.0f};
  constexpr float lambda = 0.5f;
  constexpr float alpha = 0.2f;
  constexpr float beta = 0.8f;
  constexpr float epsilon = 1e-6f;
  constexpr float max_norm = 1.0f;
  constexpr int64_t step = 0;
  constexpr float loss_scale = 1.f;
  constexpr float scaled_g_norm = 1.f;

  // Intermediate and output buffers of the optimizer.
  std::vector<float> m_new = {0.0f};
  std::vector<float> v_new = {0.0f};
  std::vector<float> w_new = {0.0f};

  run_multi_tensor_lamb_test(
      {shape},
      eta,
      {w},
      {g},
      {m},
      {v},
      {lambda},
      {alpha},
      {beta},
      {epsilon},
      {max_norm},
      step,
      loss_scale,
      &scaled_g_norm);
}

// A optimizer test with an 1-element tensor.
TEST(OptimizerTest, LambOptimizerTestScalar_NonDefaultMaxNormClipping) {
  // Input tensors and attributes.
  const std::vector<int64_t> shape = {(int64_t)1};
  constexpr float eta = 0.5f;
  const std::vector<float> w = {1.0f};
  const std::vector<float> g = {3.0f};
  const std::vector<float> m = {-10.0f};
  const std::vector<float> v = {1.0f};
  constexpr float lambda = 0.5f;
  constexpr float alpha = 0.2f;
  constexpr float beta = 0.8f;
  constexpr float epsilon = 1e-6f;
  constexpr float max_norm = 0.1f;

  // Intermediate and output buffers of the optimizer.
  std::vector<float> m_new = {0.0f};
  std::vector<float> v_new = {0.0f};
  std::vector<float> w_new = {0.0f};

  constexpr int64_t step = 0;
  constexpr float loss_scale = 1.0f;
  const float scaled_g_norm = g[0];
  run_multi_tensor_lamb_test(
      {shape},
      eta,
      {w},
      {g},
      {m},
      {v},
      {lambda},
      {alpha},
      {beta},
      {epsilon},
      {max_norm},
      step,
      loss_scale,
      &scaled_g_norm);
}

TEST(OptimizerTest, LambOptimizerTestScalarScaling) {
  // Input tensors and attributes.
  const std::vector<int64_t> shape = {(int64_t)1};
  constexpr float eta = 0.5f;
  const std::vector<float> w = {1.0f};
  const std::vector<float> g = {3.0f};
  const std::vector<float> m = {-10.0f};
  const std::vector<float> v = {1.0f};
  constexpr float lambda = 0.5f;
  constexpr float alpha = 0.2f;
  constexpr float beta = 0.8f;
  constexpr float epsilon = 1e-6f;
  constexpr float max_norm = 1.0f;
  constexpr int64_t step = 0;
  constexpr float loss_scale = 8.f;
  constexpr float scaled_g_norm = 4.f;

  // Intermediate and output buffers of the optimizer.
  std::vector<float> m_new = {0.0f};
  std::vector<float> v_new = {0.0f};
  std::vector<float> w_new = {0.0f};

  run_multi_tensor_lamb_test(
      {shape},
      eta,
      {w},
      {g},
      {m},
      {v},
      {lambda},
      {alpha},
      {beta},
      {epsilon},
      {max_norm},
      step,
      loss_scale,
      &scaled_g_norm);
}

TEST(OptimizerTest, LambOptimizerTestExternalBaseline) {
  // Input tensors and attributes.
  const std::vector<int64_t> shape = {2, 5};
  const std::vector<float> eta = {0.1f};
  const std::vector<float> w = {
      0.01379026f, 0.15308191f, -0.24356517f, -0.21798165f, -0.13770047f, 0.09694599f,
      -0.02223516f, 0.2664228f, -0.01177993f, 0.06832688f};
  const std::vector<float> g = {
      -6.048543f, 10.569487f, -9.207029f, -0.57407373f,
      5.884985f, -0.21047728f, 3.539946f, -5.957566f, -9.343748f, 1.1502024f};
  const std::vector<float> m = {
      -5.9078765f, 9.673933f, -8.731428f, -0.6227454f, 5.284312f, -0.27138948f,
      3.443532f, -5.681713f, -8.72421f, 1.1441823f};
  const std::vector<float> v = {
      4.2659229e+01f, 1.1438165e+02f, 9.3179581e+01f, 4.7399229e-01f, 3.4129276e+01f,
      9.0019435e-02f, 1.4493006e+01f, 3.9455612e+01f, 9.3025581e+01f, 1.6000764e+0f};

  constexpr float lambda = 0.1f;
  constexpr float alpha = 0.1f;
  constexpr float beta = 0.01f;
  constexpr float epsilon = 0.1f;
  constexpr float max_norm = 1.0f;

  std::vector<float> w_new = {
      0.02979828f, 0.13677707f, -0.22708717f, -0.20361158f, -0.15338624f, 0.1081504f,
      -0.03804127f, 0.28198114f, 0.00430069f, 0.05319814f};
  std::vector<float> g_new = {
      0.01600802f, -0.01630484f, 0.01647800f, 0.01437007f, -0.01568577f, 0.01120441f,
      -0.01580611f, 0.01555834f, 0.01608062f, -0.01512874f};
  std::vector<float> m_new = {
      -6.0344763f, 10.479931f, -9.15947f, -0.57894087f, 5.824918f, -0.2165685f,
      3.5303047f, -5.9299808f, -9.281795f, 1.1496004f};
  std::vector<float> v_new = {
      3.6645618e+01f, 1.1174072e+02f, 8.4853485e+01f, 3.3100498e-01f, 3.4628010e+01f,
      4.4757873e-02f, 1.2550836e+01f, 3.5532223e+01f, 8.7362823e+01f, 1.3257366e+00f};

  // Output new weights
  run_lamb_test_with_baseline(
      shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, max_norm, w_new, {}, m_new, v_new);

  // Output new gradients
  run_lamb_test_with_baseline(
      shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, max_norm, {}, g_new, m_new, v_new);
}

TEST(OptimizerTest, LambOptimizerTestExternalBaselineDouble) {
  // Input tensors and attributes.
  const std::vector<int64_t> shape = {2, 5};
  const std::vector<double> eta = {0.1f};
  const std::vector<double> w = {
      0.01379026, 0.15308191, -0.24356517, -0.21798165, -0.13770047, 0.09694599,
      -0.02223516, 0.2664228, -0.01177993, 0.06832688};
  const std::vector<double> g = {
      -6.048543, 10.569487, -9.207029, -0.57407373,
      5.884985, -0.21047728, 3.539946, -5.957566, -9.343748, 1.1502024};
  const std::vector<double> m = {
      -5.9078765, 9.673933, -8.731428, -0.6227454, 5.284312, -0.27138948,
      3.443532, -5.681713, -8.72421, 1.1441823};
  const std::vector<double> v = {
      4.2659229e+01, 1.1438165e+02, 9.3179581e+01, 4.7399229e-01, 3.4129276e+01,
      9.0019435e-02, 1.4493006e+01, 3.9455612e+01, 9.3025581e+01, 1.6000764e+0};

  constexpr float lambda = 0.1f;
  constexpr float alpha = 0.1f;
  constexpr float beta = 0.01f;
  constexpr float epsilon = 0.1f;
  constexpr float max_norm = 1.0f;

  std::vector<double> w_new = {
      0.02979828, 0.13677707, -0.22708717, -0.20361158, -0.15338624, 0.1081504,
      -0.03804127, 0.28198114, 0.00430069, 0.05319814};
  std::vector<double> g_new = {
      0.01600802, -0.01630484, 0.016478, 0.01437007, -0.01568577, 0.01120441,
      -0.01580611, 0.01555834, 0.01608062, -0.01512874};
  std::vector<double> m_new = {
      -6.0344763, 10.479931, -9.15947, -0.57894087, 5.824918, -0.2165685,
      3.5303047, -5.9299808, -9.281795, 1.1496004};
  std::vector<double> v_new = {
      3.6645618e+01, 1.1174072e+02, 8.4853485e+01, 3.3100498e-01, 3.4628010e+01,
      4.4757873e-02, 1.2550836e+01, 3.5532223e+01, 8.7362823e+01, 1.3257366e+00};

  // Output new weights
  run_lamb_test_with_baseline(
      shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, max_norm, w_new, {}, m_new, v_new);

  // Output new gradients
  run_lamb_test_with_baseline(
      shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, max_norm, {}, g_new, m_new, v_new);
}

TEST(OptimizerTest, LambOptimizerTest5DTensorMixPrecision32_16) {
  const std::vector<int64_t> shape = {2, 2, 2, 1, 1};
  const std::vector<float> eta = {0.5f};
  const std::vector<float> w = {1.0f, 2.0f, 2.5f, 1.5f, 1.0f, 2.0f, 2.0f, 1.5f};
  const std::vector<float> g = {-1.0f, 1.0f, -1.0f, 1.0f, -1.0f, 1.0f, -1.0f, 0.8f};
  const std::vector<float> m = {1.0f, 2.0f, -0.25f, 1.1f, 1.0f, 2.0f, -0.21f, 1.1f};
  const std::vector<float> v = {1.5f, 1.0f, 1.1f, 0.76f, 1.5f, 1.0f, 1.5f, 0.76f};

  constexpr float lambda = 1.5f;
  constexpr float alpha = 1.5f;
  constexpr float beta = 1.5f;
  constexpr float epsilon = 1.0f;
  constexpr float max_norm = 1.0f;
  constexpr float loss_scale = 1.0f;
  run_lamb_mix_precision_test(
      shape, eta, w, g, m, v, lambda, alpha, beta, epsilon, max_norm);

  float gradient_norm = GetGradientL2Norm(g);
  // gradient clipping
  run_lamb_mix_precision_test(
      shape, eta, w, g, m, v,
      lambda, alpha, beta, epsilon, max_norm, 0, loss_scale, &gradient_norm);
}

TEST(OptimizerTest, LambOptimizerTestSimpleBaselineMixPrecision32_16) {
  const std::vector<int64_t> shape = {2, 1};
  const std::vector<float> eta = {1.0f};
  const std::vector<float> w = {1.0f, 1.0f};
  const std::vector<float> g = {-1.0f, 1.0f};
  const std::vector<float> m = {1.0f, 1.0f};
  const std::vector<float> v = {0.0f, 0.0f};

  constexpr float lambda = 0.0f;
  constexpr float alpha = 1.0f;
  constexpr float beta = 1.0f;
  constexpr float epsilon = 1.0f;
  constexpr float max_norm = 1.0f;
  constexpr float loss_scale = 1.0f;
  run_lamb_mix_precision_test(
      shape, eta, w, g, m, v,
      lambda, alpha, beta, epsilon, max_norm);

  // gradient clipping
  float gradient_norm = GetGradientL2Norm(g);
  run_lamb_mix_precision_test(
      shape, eta, w, g, m, v,
      lambda, alpha, beta, epsilon, max_norm, 0, loss_scale, &gradient_norm);
}

TEST(OptimizerTest, LambOptimizerTestBaselineMixPrecision32_16) {
  const std::vector<int64_t> shape = {2, 1};
  const std::vector<float> eta = {0.1f};
  const std::vector<float> w = {-1.5f, 2.4f};
  const std::vector<float> g = {-0.75f, 1.2f};
  const std::vector<float> m = {0.87f, -0.94f};
  const std::vector<float> v = {0.12f, 0.28f};

  constexpr float lambda = 0.25f;
  constexpr float alpha = 0.9f;
  constexpr float beta = 0.95f;
  constexpr float epsilon = 0.33f;
  constexpr float max_norm = 1.0f;
  constexpr float loss_scale = 1.0f;

  run_lamb_mix_precision_test(
      shape, eta, w, g, m, v,
      lambda, alpha, beta, epsilon, max_norm);

  // gradient clipping
  float gradient_norm = GetGradientL2Norm(g);
  run_lamb_mix_precision_test(
      shape, eta, w, g, m, v,
      lambda, alpha, beta, epsilon, max_norm, 0, loss_scale, &gradient_norm);
}

TEST(OptimizerTest, LambOptimizerTestScalarMixPrecision32_16) {
  const std::vector<int64_t> shape = {1};
  const std::vector<float> eta = {0.1f};
  const std::vector<float> w = {-1.5f};
  const std::vector<float> g = {-0.75f};
  const std::vector<float> m = {0.87f};
  const std::vector<float> v = {0.12f};

  constexpr float lambda = 0.25f;
  constexpr float alpha = 0.9f;
  constexpr float beta = 0.95f;
  constexpr float epsilon = 0.33f;
  constexpr float max_norm = 1.0f;
  constexpr float loss_scale = 1.0f;

  run_lamb_mix_precision_test(
      shape, eta, w, g, m, v,
      lambda, alpha, beta, epsilon, max_norm);

  // gradient clipping
  float gradient_norm = GetGradientL2Norm(g);
  run_lamb_mix_precision_test(
      shape, eta, w, g, m, v,
      lambda, alpha, beta, epsilon, max_norm, 2, loss_scale, &gradient_norm);
}

TEST(OptimizerTest, LambOptimizerTestScalarMixPrecision32_16_NoDefaultMaxNormClipping) {
  const std::vector<int64_t> shape = {1};
  const std::vector<float> eta = {0.1f};
  const std::vector<float> w = {-1.5f};
  const std::vector<float> g = {-0.75f};
  const std::vector<float> m = {0.87f};
  const std::vector<float> v = {0.12f};

  constexpr float lambda = 0.25f;
  constexpr float alpha = 0.9f;
  constexpr float beta = 0.95f;
  constexpr float epsilon = 0.33f;
  constexpr float max_norm = 0.1f;
  constexpr float loss_scale = 1.0f;

  run_lamb_mix_precision_test(
      shape, eta, w, g, m, v,
      lambda, alpha, beta, epsilon, max_norm, 2);

  // gradient clipping
  float gradient_norm = GetGradientL2Norm(g);
  run_lamb_mix_precision_test(
      shape, eta, w, g, m, v,
      lambda, alpha, beta, epsilon, max_norm, 2, loss_scale, &gradient_norm);
}

TEST(OptimizerTest, LambOptimizerTestLarge) {
  // Input tensors and attributes.
  for (const auto& size : {55667, 1944006, 3907584}) {
    const std::vector<int64_t> shape = {static_cast<int64_t>(size)};
    constexpr float eta = 0.5f;
    std::vector<float> w(size);
    std::vector<float> g(size);
    std::vector<float> m(size);
    std::vector<float> v(size);

    std::random_device random_device;
    std::mt19937 random_engine(0);
    std::uniform_real_distribution<float> dist(0.1f, 1.0f);
    for (int i = 0; i < size; ++i) {
      w[i] = dist(random_engine);
      g[i] = dist(random_engine);
      m[i] = dist(random_engine);
      v[i] = dist(random_engine);
    }

    constexpr float lambda = 0.5f;
    constexpr float alpha = 0.2f;
    constexpr float beta = 0.8f;
    constexpr float epsilon = 1e-6f;
    constexpr float max_norm = 1.0f;
    constexpr int64_t step = 0;
    constexpr float loss_scale = 1.f;
    constexpr float scaled_g_norm = 1.f;

    run_multi_tensor_lamb_test(
        {shape},
        eta,
        {w},
        {g},
        {m},
        {v},
        {lambda},
        {alpha},
        {beta},
        {epsilon},
        {max_norm},
        step,
        loss_scale,
        &scaled_g_norm);
  }
}

TEST(OptimizerTest, LambOptimizerMultiTensorRatio) {
  constexpr int group_count = 127;
  std::random_device random_device;
  std::mt19937 random_engine(0);
  std::uniform_real_distribution<float> dist(0.1f, 1.0f);
  std::uniform_int_distribution<int64_t> dist_int(1, 1228);

  std::vector<int64_t> sizes(group_count);
  std::vector<std::vector<int64_t>> shapes(group_count);

  std::vector<std::vector<float>> ws(group_count);
  std::vector<std::vector<float>> gs(group_count);
  std::vector<std::vector<float>> ms(group_count);
  std::vector<std::vector<float>> vs(group_count);
  std::vector<float> alphas(group_count);
  std::vector<float> betas(group_count);
  std::vector<float> lambdas(group_count);
  std::vector<float> epsilons(group_count);
  std::vector<float> max_norms(group_count);

  const float eta = dist(random_engine);

  for (int64_t i = 0; i < group_count; ++i) {
    const auto size = dist_int(random_engine);
    sizes[i] = size;
    shapes[i] = std::vector<int64_t>(1, size);

    ws[i] = std::vector<float>(sizes[i]);
    gs[i] = std::vector<float>(sizes[i]);
    ms[i] = std::vector<float>(sizes[i]);
    vs[i] = std::vector<float>(sizes[i]);

    for (int64_t j = 0; j < sizes[i]; ++j) {
      ws[i][j] = dist(random_engine);
      gs[i][j] = dist(random_engine);
      ms[i][j] = dist(random_engine);
      vs[i][j] = dist(random_engine);
    }

    alphas[i] = dist(random_engine);
    betas[i] = dist(random_engine);
    lambdas[i] = dist(random_engine);
    epsilons[i] = dist(random_engine);
    max_norms[i] = dist(random_engine);
  }

  constexpr int64_t step = 0;
  float loss_scale = 1.f;
  constexpr float scaled_g_norm = 1.f;

  run_multi_tensor_lamb_test(
      shapes, eta,
      ws, gs, ms, vs,
      lambdas, alphas, betas, epsilons, max_norms,
      step, loss_scale, &scaled_g_norm, 0.3f, 0.7f);

  run_multi_tensor_lamb_test(
      shapes, eta,
      ws, gs, ms, vs,
      lambdas, alphas, betas, epsilons, max_norms,
      step, loss_scale, &scaled_g_norm);
}
#endif
}  // namespace
}  // namespace test
}  // namespace onnxruntime
