#include <gtest/gtest.h>

#include <torch/torch.h>

#include <test/cpp/api/support.h>

using namespace torch::nn;

struct TransformerTest : torch::test::SeedingFixture {};

// a generic function to set constants for parameters so we have fixed result
// for deterministic test
template <typename Model>
void set_parameter_to_constants(
    Model& model,
    const torch::TensorOptions& tensor_options) {
  torch::NoGradGuard guard;
  for (auto& p : model->parameters()) {
    auto sz = p.view(-1).size(0);
    p.copy_(torch::cos(torch::arange(0, sz, tensor_options).view(p.sizes())));
  }
}

// a generic function to provide consistent encoder/decoder layer for all the
// transformer tests
template <typename T_LAYER, typename T_OPTIONS>
T_LAYER get_a_test_layer(
    const torch::TensorOptions& tensor_options,
    bool use_callable_activation) {
  int64_t d_model = 4;
  int64_t nhead = 2;
  int64_t dim_feedforward = 16;
  double dropout = 0.0;

  // activation is always ReLU here and it can be adjusted later depending on
  // the usage
  T_LAYER layer(T_OPTIONS(d_model, nhead)
                    .dim_feedforward(dim_feedforward)
                    .dropout(dropout));
  if (tensor_options.device() == torch::kCUDA) {
    layer->to(torch::kCUDA);
  }
  if (use_callable_activation) {
    layer.get()->options.activation(
        [&](const torch::Tensor& t) { return torch::nn::functional::relu(t); });
  }

  // set constant weights of the model
  set_parameter_to_constants<T_LAYER>(layer, tensor_options);

  return layer;
}

void transformer_encoder_layer_test_helper(
    bool is_cuda,
    bool use_callable_activation) {
  // this is a deterministic test for TransformerEncoderLayer
  torch::Device device = is_cuda ? torch::kCUDA : torch::kCPU;
  torch::TensorOptions tensor_options =
      torch::TensorOptions().dtype(torch::kFloat32).device(device);

  TransformerEncoderLayer model =
      get_a_test_layer<TransformerEncoderLayer, TransformerEncoderLayerOptions>(
          tensor_options, use_callable_activation);

  // relu test case 1
  torch::Tensor encoder_input =
      torch::tensor({{{20, 30, 40, 50}}}, tensor_options);
  torch::Tensor result = model(encoder_input).detach();
  torch::Tensor ref_output = torch::tensor(
      {{{2.258703, 0.127985, -0.697881, 0.170862}}}, tensor_options);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  // all 0 values are NOT masked. This should't mask anything
  torch::Tensor mask = torch::tensor({{0}}, tensor_options) == 1;
  result = model(
               encoder_input,
               /*src_mask=*/torch::Tensor{},
               /*src_key_padding_mask=*/mask)
               .detach();
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  // all 1 values are masked. Since there is only 1 input embedding this will
  // result in nan.
  mask = torch::tensor({{1}}, tensor_options) == 1;
  result = model(
               encoder_input,
               /*src_mask=*/torch::Tensor{},
               /*src_key_padding_mask=*/mask)
               .detach();
  ASSERT_TRUE(torch::isnan(result).all().item().to<bool>());

  // relu test case 2
  encoder_input =
      torch::tensor({{{1, 2, 3, 4}}, {{5, 6, 7, 8}}}, tensor_options);
  result = model(encoder_input).detach();
  ref_output = torch::tensor(
      {{{2.272644, 0.119035, -0.691669, 0.153486}},
       {{2.272644, 0.119035, -0.691669, 0.153486}}},
      tensor_options);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  // all 0 values are NOT masked
  mask = torch::tensor({{0, 0}}, tensor_options) == 1;
  result = model(
               encoder_input,
               /*src_mask=*/torch::Tensor{},
               /*src_key_padding_mask=*/mask)
               .detach();
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  // mask with 1 and 0
  mask = torch::tensor({{1, 0}}, tensor_options) == 1;
  result = model(
               encoder_input,
               /*src_mask=*/torch::Tensor{},
               /*src_key_padding_mask=*/mask)
               .detach();
  ref_output = torch::tensor(
      {{{2.301516, 0.092249, -0.679101, 0.103088}},
       {{2.301516, 0.092249, -0.679101, 0.103088}}},
      tensor_options);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  // relu test case 3
  encoder_input = torch::tensor(
      {{{0.7462, 0.6653, 0.5679, 0.4891}, {0.5387, 0.1655, 0.3565, 0.0471}},
       {{0.8335, 0.2799, 0.5031, 0.2947}, {0.1402, 0.0318, 0.7636, 0.1346}},
       {{0.6333, 0.9344, 0.1376, 0.9938}, {0.8924, 0.2872, 0.6692, 0.2944}},
       {{0.9897, 0.6915, 0.3154, 0.1733}, {0.8645, 0.3513, 0.3064, 0.0767}},
       {{0.8117, 0.2366, 0.4838, 0.7881}, {0.3718, 0.4945, 0.9511, 0.0864}}},
      tensor_options);
  result = model(encoder_input).detach();
  ref_output = torch::tensor(
      {{{2.428589, 0.020835, -0.602055, -0.085249},
        {2.427987, 0.021213, -0.602496, -0.084103}},
       {{2.424689, 0.019155, -0.604793, -0.085672},
        {2.413863, 0.022211, -0.612486, -0.072490}},
       {{2.433774, 0.021598, -0.598343, -0.087548},
        {2.425104, 0.019748, -0.604515, -0.084839}},
       {{2.436185, 0.022682, -0.596625, -0.087261},
        {2.433556, 0.021891, -0.598509, -0.086832}},
       {{2.416246, 0.017512, -0.610712, -0.082961},
        {2.422901, 0.024187, -0.606178, -0.074929}}},
      tensor_options);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  // all 0 values are NOT masked
  mask = torch::zeros({2, 5}, tensor_options) == 1;
  result = model(
               encoder_input,
               /*src_mask=*/torch::Tensor{},
               /*src_key_padding_mask=*/mask)
               .detach();
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  // mask with 0s and 1s
  mask[0][1] = 1;
  mask[1][3] = 1;
  mask[1][4] = 1;
  result = model(
               encoder_input,
               /*src_mask=*/torch::Tensor{},
               /*src_key_padding_mask=*/mask)
               .detach();
  ref_output = torch::tensor(
      {{{2.429026, 0.020793, -0.601741, -0.085642},
        {2.428811, 0.021445, -0.601912, -0.084252}},
       {{2.425009, 0.019155, -0.604566, -0.085899},
        {2.415408, 0.02249, -0.611415, -0.073}},
       {{2.434199, 0.021682, -0.598039, -0.087699},
        {2.42598, 0.019941, -0.603896, -0.085091}},
       {{2.436457, 0.022736, -0.59643, -0.08736},
        {2.434021, 0.022093, -0.598179, -0.08679}},
       {{2.416531, 0.017498, -0.610513, -0.083181},
        {2.4242, 0.024653, -0.605266, -0.074959}}},
      tensor_options);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  // gelu test case 1
  model.get()->options.activation(torch::kGELU);
  encoder_input = torch::tensor({{{20, 30, 40, 50}}}, tensor_options);
  result = model(encoder_input).detach();
  ref_output = torch::tensor(
      {{{2.249815, 0.131006, -0.702199, 0.177868}}}, tensor_options);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  // gelu test case 2
  encoder_input = torch::tensor(
      {{{0.7462, 0.6653, 0.5679, 0.4891}, {0.5387, 0.1655, 0.3565, 0.0471}},
       {{0.8335, 0.2799, 0.5031, 0.2947}, {0.1402, 0.0318, 0.7636, 0.1346}},
       {{0.6333, 0.9344, 0.1376, 0.9938}, {0.8924, 0.2872, 0.6692, 0.2944}},
       {{0.9897, 0.6915, 0.3154, 0.1733}, {0.8645, 0.3513, 0.3064, 0.0767}},
       {{0.8117, 0.2366, 0.4838, 0.7881}, {0.3718, 0.4945, 0.9511, 0.0864}}},
      tensor_options);
  result = model(encoder_input);
  ref_output = torch::tensor(
      {{{2.42163188, 0.03227153, -0.60714219, -0.05908082},
        {2.42151276, 0.03302179, -0.60722523, -0.05762651}},
       {{2.41926761, 0.02974034, -0.60879519, -0.0621269},
        {2.41626395, 0.03539356, -0.61087842, -0.04978623}},
       {{2.42382808, 0.03218872, -0.6055963, -0.06073591},
        {2.41983477, 0.03085259, -0.60840145, -0.06046414}},
       {{2.42500749, 0.03328855, -0.60476388, -0.0595334},
        {2.4237977, 0.03290575, -0.60561789, -0.05940082}},
       {{2.41383916, 0.02686345, -0.61256377, -0.06380707},
        {2.42000277, 0.03800944, -0.60824798, -0.04754947}}},
      tensor_options);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));
}

TEST_F(TransformerTest, TransformerEncoderLayer) {
  transformer_encoder_layer_test_helper(
      /*is_cuda=*/false, /*use_callable_activation=*/false);
  transformer_encoder_layer_test_helper(
      /*is_cuda=*/false, /*use_callable_activation=*/true);
}

TEST_F(TransformerTest, TransformerEncoderLayer_CUDA) {
  transformer_encoder_layer_test_helper(
      /*is_cuda=*/true, /*use_callable_activation=*/false);
  transformer_encoder_layer_test_helper(
      /*is_cuda=*/true, /*use_callable_activation=*/true);
}

void transformer_decoder_layer_test_helper(
    bool is_cuda,
    bool use_callable_activation) {
  torch::Device device = is_cuda ? torch::kCUDA : torch::kCPU;
  torch::TensorOptions tensor_options =
      torch::TensorOptions().dtype(torch::kFloat32).device(device);

  TransformerDecoderLayer model =
      get_a_test_layer<TransformerDecoderLayer, TransformerDecoderLayerOptions>(
          tensor_options, use_callable_activation);

  // deterministic input
  torch::Tensor decoder_input =
      torch::tensor({{{20, 30, 40, 50}}}, tensor_options);
  torch::Tensor memory_input =
      torch::tensor({{{60, 70, 80, 90}}}, tensor_options);
  torch::Tensor result = model(decoder_input, memory_input).detach();
  torch::Tensor ref_output = torch::tensor(
      {{{2.314351, 0.094805, -0.671322, 0.101977}}}, tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // deterministic input
  decoder_input =
      torch::tensor({{{9, 10, 11, 12}}, {{11, 12, 13, 14}}}, tensor_options);
  memory_input = torch::tensor({{{1, 2, 3, 4}}}, tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.422245, 0.051716, -0.606338, -0.024756}},
       {{2.422245, 0.051716, -0.606338, -0.024756}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // deterministic input
  decoder_input =
      torch::tensor({{{1, 2, 3, 4}}, {{5, 6, 7, 8}}}, tensor_options);
  memory_input =
      torch::tensor({{{9, 10, 11, 12}}, {{11, 12, 13, 14}}}, tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.343536, 0.085561, -0.654954, 0.074991}},
       {{2.343536, 0.085561, -0.654954, 0.074991}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // deterministic input
  decoder_input = torch::tensor(
      {{{0.4517, 0.6793, 0.5313, 0.0034}, {0.2678, 0.3677, 0.4459, 0.7166}},
       {{0.8100, 0.3716, 0.4096, 0.1976}, {0.6958, 0.8844, 0.6081, 0.8315}},
       {{0.0494, 0.9343, 0.5955, 0.3830}, {0.5404, 0.3464, 0.9378, 0.6200}}},
      tensor_options);
  memory_input = torch::tensor(
      {{{0.7462, 0.6653, 0.5679, 0.4891}, {0.5387, 0.1655, 0.3565, 0.0471}},
       {{0.8335, 0.2799, 0.5031, 0.2947}, {0.1402, 0.0318, 0.7636, 0.1346}},
       {{0.6333, 0.9344, 0.1376, 0.9938}, {0.8924, 0.2872, 0.6692, 0.2944}},
       {{0.9897, 0.6915, 0.3154, 0.1733}, {0.8645, 0.3513, 0.3064, 0.0767}},
       {{0.8117, 0.2366, 0.4838, 0.7881}, {0.3718, 0.4945, 0.9511, 0.0864}}},
      tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.430065, 0.027862, -0.601136, -0.073096},
        {2.431935, 0.028907, -0.599809, -0.072488}},
       {{2.428457, 0.027053, -0.602275, -0.073462},
        {2.431970, 0.029387, -0.599789, -0.071621}},
       {{2.431934, 0.028196, -0.599802, -0.073809},
        {2.432306, 0.028858, -0.599542, -0.072846}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // key_padding_mask
  torch::Tensor t_mask = {};
  torch::Tensor m_mask = {};
  torch::Tensor key_padding_mask = torch::zeros({2, 3}, tensor_options) == 1;
  result = model(decoder_input, memory_input, t_mask, m_mask, key_padding_mask)
               .detach();
  ref_output = torch::tensor(
      {{{2.430065, 0.027862, -0.601136, -0.073096},
        {2.431935, 0.028907, -0.599809, -0.072488}},
       {{2.428457, 0.027053, -0.602275, -0.073462},
        {2.431970, 0.029387, -0.599789, -0.071621}},
       {{2.431934, 0.028196, -0.599802, -0.073809},
        {2.432306, 0.028858, -0.599542, -0.072846}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // key_padding_mask
  key_padding_mask[0][2] = 1;
  key_padding_mask[1][1] = 1;
  key_padding_mask[1][2] = 1;
  result = model(decoder_input, memory_input, t_mask, m_mask, key_padding_mask)
               .detach();
  ref_output = torch::tensor(
      {{{2.430025, 0.027643, -0.601164, -0.073476},
        {2.4323, 0.029375, -0.599553, -0.071881}},
       {{2.428523, 0.026838, -0.602226, -0.07391},
        {2.432634, 0.029842, -0.599318, -0.071253}},
       {{2.432278, 0.028152, -0.599555, -0.074139},
        {2.432659, 0.029244, -0.599294, -0.072382}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // memory_key_padding_mask
  torch::Tensor t_key_padding_mask = {};
  key_padding_mask = torch::zeros({2, 5}, tensor_options) == 1;
  result = model(
               decoder_input,
               memory_input,
               t_mask,
               m_mask,
               t_key_padding_mask,
               key_padding_mask)
               .detach();
  ref_output = torch::tensor(
      {{{2.430065, 0.027862, -0.601136, -0.073096},
        {2.431935, 0.028907, -0.599809, -0.072488}},
       {{2.428457, 0.027053, -0.602275, -0.073462},
        {2.431970, 0.029387, -0.599789, -0.071621}},
       {{2.431934, 0.028196, -0.599802, -0.073809},
        {2.432306, 0.028858, -0.599542, -0.072846}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // memory_key_padding_mask
  key_padding_mask[0][4] = 1;
  key_padding_mask[1][3] = 1;
  key_padding_mask[1][4] = 1;
  result = model(
               decoder_input,
               memory_input,
               t_mask,
               m_mask,
               t_key_padding_mask,
               key_padding_mask)
               .detach();
  ref_output = torch::tensor(
      {{{2.429757, 0.027358, -0.601351, -0.073816},
        {2.432692, 0.028583, -0.599263, -0.073634}},
       {{2.428247, 0.02662, -0.602419, -0.074123},
        {2.432657, 0.029055, -0.599293, -0.072732}},
       {{2.431515, 0.027687, -0.600096, -0.074459},
        {2.433075, 0.028543, -0.598987, -0.073985}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));
}

TEST_F(TransformerTest, TransformerDecoderLayer) {
  transformer_decoder_layer_test_helper(
      /*is_cuda=*/false, /*use_callable_activation=*/false);
  transformer_decoder_layer_test_helper(
      /*is_cuda=*/false, /*use_callable_activation=*/true);
}

TEST_F(TransformerTest, TransformerDecoderLayer_CUDA) {
  transformer_decoder_layer_test_helper(
      /*is_cuda=*/true, /*use_callable_activation=*/false);
  transformer_decoder_layer_test_helper(
      /*is_cuda=*/true, /*use_callable_activation=*/true);
}

void transformer_decoder_layer_test_helper_gelu(
    bool is_cuda,
    bool use_callable_activation) {
  torch::Device device = is_cuda ? torch::kCUDA : torch::kCPU;
  torch::TensorOptions tensor_options =
      torch::TensorOptions().dtype(torch::kFloat32).device(device);

  TransformerDecoderLayer model =
      get_a_test_layer<TransformerDecoderLayer, TransformerDecoderLayerOptions>(
          tensor_options, use_callable_activation);
  if (use_callable_activation) {
    model.get()->options.activation(
        [&](const torch::Tensor& t) { return torch::nn::functional::gelu(t); });
  } else {
    model.get()->options.activation(torch::kGELU);
  }

  // deterministic input
  torch::Tensor decoder_input =
      torch::tensor({{{20, 30, 40, 50}}}, tensor_options);
  torch::Tensor memory_input =
      torch::tensor({{{60, 70, 80, 90}}}, tensor_options);
  torch::Tensor result = model(decoder_input, memory_input).detach();
  torch::Tensor ref_output = torch::tensor(
      {{{2.306435, 0.095946, -0.675796, 0.10687}}}, tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // deterministic input
  decoder_input =
      torch::tensor({{{9, 10, 11, 12}}, {{11, 12, 13, 14}}}, tensor_options);
  memory_input = torch::tensor({{{1, 2, 3, 4}}}, tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.415448, 0.054389, -0.610932, -0.0156613}},
       {{2.415448, 0.054389, -0.610932, -0.0156613}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // deterministic input
  decoder_input =
      torch::tensor({{{1, 2, 3, 4}}, {{5, 6, 7, 8}}}, tensor_options);
  memory_input =
      torch::tensor({{{9, 10, 11, 12}}, {{11, 12, 13, 14}}}, tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.338531, 0.087709, -0.65776, 0.080646}},
       {{2.338531, 0.087709, -0.65776, 0.080646}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // deterministic input
  decoder_input = torch::tensor(
      {{{0.4517, 0.6793, 0.5313, 0.0034}, {0.2678, 0.3677, 0.4459, 0.7166}},
       {{0.8100, 0.3716, 0.4096, 0.1976}, {0.6958, 0.8844, 0.6081, 0.8315}},
       {{0.0494, 0.9343, 0.5955, 0.3830}, {0.5404, 0.3464, 0.9378, 0.6200}}},
      tensor_options);
  memory_input = torch::tensor(
      {{{0.7462, 0.6653, 0.5679, 0.4891}, {0.5387, 0.1655, 0.3565, 0.0471}},
       {{0.8335, 0.2799, 0.5031, 0.2947}, {0.1402, 0.0318, 0.7636, 0.1346}},
       {{0.6333, 0.9344, 0.1376, 0.9938}, {0.8924, 0.2872, 0.6692, 0.2944}},
       {{0.9897, 0.6915, 0.3154, 0.1733}, {0.8645, 0.3513, 0.3064, 0.0767}},
       {{0.8117, 0.2366, 0.4838, 0.7881}, {0.3718, 0.4945, 0.9511, 0.0864}}},
      tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.42049104, 0.03443088, -0.60793706, -0.05436271},
        {2.42210631, 0.03546578, -0.60679895, -0.05357488}},
       {{2.41907674, 0.0336104, -0.60892977, -0.05490462},
        {2.42216881, 0.03586554, -0.6067524, -0.05289126}},
       {{2.42205716, 0.03488046, -0.60683681, -0.05460596},
        {2.42240309, 0.0354595, -0.60659063, -0.05378816}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));
}

TEST_F(TransformerTest, TransformerDecoderLayer_gelu) {
  transformer_decoder_layer_test_helper_gelu(
      /*is_cuda=*/false, /*use_callable_activation=*/false);
  transformer_decoder_layer_test_helper_gelu(
      /*is_cuda=*/false, /*use_callable_activation=*/true);
}

TEST_F(TransformerTest, TransformerDecoderLayer_gelu_CUDA) {
  transformer_decoder_layer_test_helper_gelu(
      /*is_cuda=*/true, /*use_callable_activation=*/false);
  transformer_decoder_layer_test_helper_gelu(
      /*is_cuda=*/true, /*use_callable_activation=*/true);
}

void transformer_encoder_test_helper(
    bool is_cuda,
    bool use_callable_activation) {
  // this is a deterministic test for TransformerEncoderLayer
  torch::Device device = is_cuda ? torch::kCUDA : torch::kCPU;
  torch::TensorOptions tensor_options =
      torch::TensorOptions().dtype(torch::kFloat32).device(device);

  TransformerEncoderLayer encoder_layer =
      get_a_test_layer<TransformerEncoderLayer, TransformerEncoderLayerOptions>(
          tensor_options, use_callable_activation);

  TransformerEncoder model(TransformerEncoderOptions(encoder_layer, 1));
  if (is_cuda) {
    model->to(torch::kCUDA);
  }

  torch::Tensor encoder_input = torch::tensor(
      {{{0.7462, 0.6653, 0.5679, 0.4891}, {0.5387, 0.1655, 0.3565, 0.0471}},
       {{0.8335, 0.2799, 0.5031, 0.2947}, {0.1402, 0.0318, 0.7636, 0.1346}},
       {{0.6333, 0.9344, 0.1376, 0.9938}, {0.8924, 0.2872, 0.6692, 0.2944}},
       {{0.9897, 0.6915, 0.3154, 0.1733}, {0.8645, 0.3513, 0.3064, 0.0767}},
       {{0.8117, 0.2366, 0.4838, 0.7881}, {0.3718, 0.4945, 0.9511, 0.0864}}},
      tensor_options);
  torch::Tensor result = model(encoder_input).detach();
  torch::Tensor ref_output = torch::tensor(
      {{{2.428589, 0.020835, -0.602055, -0.085249},
        {2.427987, 0.021213, -0.602496, -0.084103}},
       {{2.424689, 0.019155, -0.604793, -0.085672},
        {2.413863, 0.022211, -0.612486, -0.072490}},
       {{2.433774, 0.021598, -0.598343, -0.087548},
        {2.425104, 0.019748, -0.604515, -0.084839}},
       {{2.436185, 0.022682, -0.596625, -0.087261},
        {2.433556, 0.021891, -0.598509, -0.086832}},
       {{2.416246, 0.017512, -0.610712, -0.082961},
        {2.422901, 0.024187, -0.606178, -0.074929}}},
      tensor_options);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  // all 0 values are NOT masked
  torch::Tensor mask = torch::zeros({2, 5}, tensor_options) == 1;
  result = model(
               encoder_input,
               /*src_mask=*/torch::Tensor{},
               /*src_key_padding_mask=*/mask)
               .detach();
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  // mask with 0s and 1s
  mask[0][1] = 1;
  mask[1][3] = 1;
  mask[1][4] = 1;
  result = model(
               encoder_input,
               /*src_mask=*/torch::Tensor{},
               /*src_key_padding_mask=*/mask)
               .detach();
  ref_output = torch::tensor(
      {{{2.429026, 0.020793, -0.601741, -0.085642},
        {2.428811, 0.021445, -0.601912, -0.084252}},
       {{2.425009, 0.019155, -0.604566, -0.085899},
        {2.415408, 0.02249, -0.611415, -0.073}},
       {{2.434199, 0.021682, -0.598039, -0.087699},
        {2.42598, 0.019941, -0.603896, -0.085091}},
       {{2.436457, 0.022736, -0.59643, -0.08736},
        {2.434021, 0.022093, -0.598179, -0.08679}},
       {{2.416531, 0.017498, -0.610513, -0.083181},
        {2.4242, 0.024653, -0.605266, -0.074959}}},
      tensor_options);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  // test case 2, multiple layers no norm
  model = TransformerEncoder(TransformerEncoderOptions(encoder_layer, 2));
  if (is_cuda) {
    model->to(torch::kCUDA);
  }
  result = model(
               encoder_input,
               /*src_mask=*/torch::Tensor{},
               /*src_key_padding_mask=*/mask)
               .detach();
  ref_output = torch::tensor(
      {{{2.419051, 0.017446, -0.608738, -0.085003},
        {2.419102, 0.017452, -0.608703, -0.085026}},
       {{2.419043, 0.017445, -0.608744, -0.084999},
        {2.419052, 0.017446, -0.608738, -0.085004}},
       {{2.419067, 0.017448, -0.608727, -0.085010},
        {2.419098, 0.017452, -0.608706, -0.085024}},
       {{2.419072, 0.017449, -0.608724, -0.085012},
        {2.419119, 0.017455, -0.608691, -0.085034}},
       {{2.419019, 0.017442, -0.608761, -0.084989},
        {2.419075, 0.017449, -0.608722, -0.085014}}},
      tensor_options);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  model = TransformerEncoder(TransformerEncoderOptions(encoder_layer, 6));
  if (is_cuda) {
    model->to(torch::kCUDA);
  }
  result = model(
               encoder_input,
               /*src_mask=*/torch::Tensor{},
               /*src_key_padding_mask=*/mask)
               .detach();
  ref_output = torch::tensor(
      {{{2.419101, 0.017453, -0.608703, -0.085025},
        {2.419101, 0.017453, -0.608704, -0.085025}},
       {{2.419101, 0.017453, -0.608703, -0.085025},
        {2.419101, 0.017453, -0.608704, -0.085025}},
       {{2.419101, 0.017453, -0.608703, -0.085025},
        {2.419101, 0.017453, -0.608704, -0.085025}},
       {{2.419101, 0.017453, -0.608703, -0.085025},
        {2.419101, 0.017453, -0.608704, -0.085025}},
       {{2.419101, 0.017453, -0.608703, -0.085025},
        {2.419101, 0.017453, -0.608704, -0.085025}}},
      tensor_options);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  // test case 3, multiple layers with norm
  LayerNorm norm(LayerNormOptions({encoder_layer.get()->options.d_model()}));
  model = TransformerEncoder(
      TransformerEncoderOptions(encoder_layer, 2).norm(AnyModule(norm)));
  if (is_cuda) {
    model->to(torch::kCUDA);
  }
  result = model(
               encoder_input,
               /*src_mask=*/torch::Tensor{},
               /*src_key_padding_mask=*/mask)
               .detach();
  ref_output = torch::tensor(
      {{{1.695949, -0.357635, -0.893077, -0.445238},
        {1.695955, -0.357639, -0.893050, -0.445266}},
       {{1.695948, -0.357634, -0.893082, -0.445233},
        {1.695950, -0.357635, -0.893077, -0.445238}},
       {{1.695951, -0.357636, -0.893069, -0.445246},
        {1.695955, -0.357639, -0.893052, -0.445264}},
       {{1.695952, -0.357636, -0.893066, -0.445249},
        {1.695957, -0.357641, -0.893041, -0.445276}},
       {{1.695946, -0.357632, -0.893095, -0.445220},
        {1.695952, -0.357637, -0.893065, -0.445251}}},
      tensor_options);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  model = TransformerEncoder(
      TransformerEncoderOptions(encoder_layer, 6).norm(AnyModule(norm)));
  if (is_cuda) {
    model->to(torch::kCUDA);
  }
  result = model(
               encoder_input,
               /*src_mask=*/torch::Tensor{},
               /*src_key_padding_mask=*/mask)
               .detach();
  ref_output = torch::tensor(
      {{{1.695955, -0.357639, -0.893051, -0.445265},
        {1.695955, -0.357639, -0.893051, -0.445265}},
       {{1.695955, -0.357639, -0.893051, -0.445265},
        {1.695955, -0.357639, -0.893051, -0.445265}},
       {{1.695955, -0.357639, -0.893051, -0.445265},
        {1.695955, -0.357639, -0.893051, -0.445265}},
       {{1.695955, -0.357639, -0.893051, -0.445265},
        {1.695955, -0.357639, -0.893051, -0.445265}},
       {{1.695955, -0.357639, -0.893051, -0.445265},
        {1.695955, -0.357639, -0.893051, -0.445265}}},
      tensor_options);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));
}

TEST_F(TransformerTest, TransformerEncoder) {
  transformer_encoder_test_helper(
      /*is_cuda=*/false, /*use_callable_activation=*/false);
  transformer_encoder_test_helper(
      /*is_cuda=*/false, /*use_callable_activation=*/true);
}

TEST_F(TransformerTest, TransformerEncoder_CUDA) {
  transformer_encoder_test_helper(
      /*is_cuda=*/true, /*use_callable_activation=*/false);
  transformer_encoder_test_helper(
      /*is_cuda=*/true, /*use_callable_activation=*/true);
}

TEST_F(TransformerTest, PrettyPrintTransformerEncoderLayer) {
  ASSERT_EQ(
      c10::str(TransformerEncoderLayer(4, 2)),
      "torch::nn::TransformerEncoderLayerImpl(\n"
      "  (self_attn): torch::nn::MultiheadAttention(\n"
      "    (out_proj): torch::nn::Linear(in_features=4, out_features=4, bias=true)\n"
      "  )\n"
      "  (linear1): torch::nn::Linear(in_features=4, out_features=2048, bias=true)\n"
      "  (dropout): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "  (linear2): torch::nn::Linear(in_features=2048, out_features=4, bias=true)\n"
      "  (norm1): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "  (norm2): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "  (dropout1): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "  (dropout2): torch::nn::Dropout(p=0.1, inplace=false)\n"
      ")");
}

TEST_F(TransformerTest, PrettyPrintTransformerEncoder) {
  LayerNorm norm = LayerNorm(LayerNormOptions({4}));
  TransformerEncoderOptions options(
      TransformerEncoderOptions(TransformerEncoderLayerOptions(4, 2), 2)
          .norm(AnyModule(norm)));
  ASSERT_EQ(
      c10::str(TransformerEncoder(options)),
      "torch::nn::TransformerEncoderImpl(\n"
      "  (layers): torch::nn::ModuleList(\n"
      "    (0): torch::nn::TransformerEncoderLayerImpl(\n"
      "      (self_attn): torch::nn::MultiheadAttention(\n"
      "        (out_proj): torch::nn::Linear(in_features=4, out_features=4, bias=true)\n"
      "      )\n"
      "      (linear1): torch::nn::Linear(in_features=4, out_features=2048, bias=true)\n"
      "      (dropout): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "      (linear2): torch::nn::Linear(in_features=2048, out_features=4, bias=true)\n"
      "      (norm1): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "      (norm2): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "      (dropout1): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "      (dropout2): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "    )\n"
      "    (1): torch::nn::TransformerEncoderLayerImpl(\n"
      "      (self_attn): torch::nn::MultiheadAttention(\n"
      "        (out_proj): torch::nn::Linear(in_features=4, out_features=4, bias=true)\n"
      "      )\n"
      "      (linear1): torch::nn::Linear(in_features=4, out_features=2048, bias=true)\n"
      "      (dropout): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "      (linear2): torch::nn::Linear(in_features=2048, out_features=4, bias=true)\n"
      "      (norm1): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "      (norm2): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "      (dropout1): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "      (dropout2): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "    )\n"
      "  )\n"
      "  (norm): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      ")");
}

TEST_F(TransformerTest, PrettyPrintTransformerDecoderLayer) {
  ASSERT_EQ(
      c10::str(TransformerDecoderLayer(4, 2)),
      "torch::nn::TransformerDecoderLayerImpl(\n"
      "  (self_attn): torch::nn::MultiheadAttention(\n"
      "    (out_proj): torch::nn::Linear(in_features=4, out_features=4, bias=true)\n"
      "  )\n"
      "  (multihead_attn): torch::nn::MultiheadAttention(\n"
      "    (out_proj): torch::nn::Linear(in_features=4, out_features=4, bias=true)\n"
      "  )\n"
      "  (linear1): torch::nn::Linear(in_features=4, out_features=2048, bias=true)\n"
      "  (dropout): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "  (linear2): torch::nn::Linear(in_features=2048, out_features=4, bias=true)\n"
      "  (norm1): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "  (norm2): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "  (norm3): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "  (dropout1): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "  (dropout2): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "  (dropout3): torch::nn::Dropout(p=0.1, inplace=false)\n"
      ")");
}

void transformer_decoder_test_helper(
    bool is_cuda,
    bool use_callable_activation) {
  // this is a deterministic test for TransformerDecoder
  torch::Device device = is_cuda ? torch::kCUDA : torch::kCPU;
  torch::TensorOptions tensor_options =
      torch::TensorOptions().dtype(torch::kFloat32).device(device);

  TransformerDecoderLayer decoder_layer =
      get_a_test_layer<TransformerDecoderLayer, TransformerDecoderLayerOptions>(
          tensor_options, use_callable_activation);

  TransformerDecoder model(TransformerDecoderOptions(decoder_layer, 1));
  if (is_cuda) {
    model->to(torch::kCUDA);
  }

  torch::Tensor decoder_input =
      torch::tensor({{{20, 30, 40, 50}}}, tensor_options);
  torch::Tensor memory_input =
      torch::tensor({{{60, 70, 80, 90}}}, tensor_options);
  torch::Tensor result = model(decoder_input, memory_input).detach();
  torch::Tensor ref_output = torch::tensor(
      {{{2.314351, 0.094805, -0.671322, 0.101977}}}, tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // deterministic input
  decoder_input =
      torch::tensor({{{9, 10, 11, 12}}, {{11, 12, 13, 14}}}, tensor_options);
  memory_input = torch::tensor({{{1, 2, 3, 4}}}, tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.422245, 0.051716, -0.606338, -0.024756}},
       {{2.422245, 0.051716, -0.606338, -0.024756}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // deterministic input
  decoder_input =
      torch::tensor({{{1, 2, 3, 4}}, {{5, 6, 7, 8}}}, tensor_options);
  memory_input =
      torch::tensor({{{9, 10, 11, 12}}, {{11, 12, 13, 14}}}, tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.343536, 0.085561, -0.654954, 0.074991}},
       {{2.343536, 0.085561, -0.654954, 0.074991}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // deterministic input
  decoder_input = torch::tensor(
      {{{0.4517, 0.6793, 0.5313, 0.0034}, {0.2678, 0.3677, 0.4459, 0.7166}},
       {{0.8100, 0.3716, 0.4096, 0.1976}, {0.6958, 0.8844, 0.6081, 0.8315}},
       {{0.0494, 0.9343, 0.5955, 0.3830}, {0.5404, 0.3464, 0.9378, 0.6200}}},
      tensor_options);
  memory_input = torch::tensor(
      {{{0.7462, 0.6653, 0.5679, 0.4891}, {0.5387, 0.1655, 0.3565, 0.0471}},
       {{0.8335, 0.2799, 0.5031, 0.2947}, {0.1402, 0.0318, 0.7636, 0.1346}},
       {{0.6333, 0.9344, 0.1376, 0.9938}, {0.8924, 0.2872, 0.6692, 0.2944}},
       {{0.9897, 0.6915, 0.3154, 0.1733}, {0.8645, 0.3513, 0.3064, 0.0767}},
       {{0.8117, 0.2366, 0.4838, 0.7881}, {0.3718, 0.4945, 0.9511, 0.0864}}},
      tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.430065, 0.027862, -0.601136, -0.073096},
        {2.431935, 0.028907, -0.599809, -0.072488}},
       {{2.428457, 0.027053, -0.602275, -0.073462},
        {2.431970, 0.029387, -0.599789, -0.071621}},
       {{2.431934, 0.028196, -0.599802, -0.073809},
        {2.432306, 0.028858, -0.599542, -0.072846}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // key_padding_mask
  torch::Tensor t_mask = {};
  torch::Tensor m_mask = {};
  torch::Tensor key_padding_mask = torch::zeros({2, 3}, tensor_options) == 1;
  result = model(decoder_input, memory_input, t_mask, m_mask, key_padding_mask)
               .detach();
  ref_output = torch::tensor(
      {{{2.430065, 0.027862, -0.601136, -0.073096},
        {2.431935, 0.028907, -0.599809, -0.072488}},
       {{2.428457, 0.027053, -0.602275, -0.073462},
        {2.431970, 0.029387, -0.599789, -0.071621}},
       {{2.431934, 0.028196, -0.599802, -0.073809},
        {2.432306, 0.028858, -0.599542, -0.072846}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // key_padding_mask
  key_padding_mask[0][2] = 1;
  key_padding_mask[1][1] = 1;
  key_padding_mask[1][2] = 1;
  result = model(decoder_input, memory_input, t_mask, m_mask, key_padding_mask)
               .detach();
  ref_output = torch::tensor(
      {{{2.430025, 0.027643, -0.601164, -0.073476},
        {2.4323, 0.029375, -0.599553, -0.071881}},
       {{2.428523, 0.026838, -0.602226, -0.07391},
        {2.432634, 0.029842, -0.599318, -0.071253}},
       {{2.432278, 0.028152, -0.599555, -0.074139},
        {2.432659, 0.029244, -0.599294, -0.072382}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // memory_key_padding_mask
  torch::Tensor t_key_padding_mask = {};
  key_padding_mask = torch::zeros({2, 5}, tensor_options) == 1;
  result = model(
               decoder_input,
               memory_input,
               t_mask,
               m_mask,
               t_key_padding_mask,
               key_padding_mask)
               .detach();
  ref_output = torch::tensor(
      {{{2.430065, 0.027862, -0.601136, -0.073096},
        {2.431935, 0.028907, -0.599809, -0.072488}},
       {{2.428457, 0.027053, -0.602275, -0.073462},
        {2.431970, 0.029387, -0.599789, -0.071621}},
       {{2.431934, 0.028196, -0.599802, -0.073809},
        {2.432306, 0.028858, -0.599542, -0.072846}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // memory_key_padding_mask
  key_padding_mask[0][4] = 1;
  key_padding_mask[1][3] = 1;
  key_padding_mask[1][4] = 1;
  result = model(
               decoder_input,
               memory_input,
               t_mask,
               m_mask,
               t_key_padding_mask,
               key_padding_mask)
               .detach();
  ref_output = torch::tensor(
      {{{2.429757, 0.027358, -0.601351, -0.073816},
        {2.432692, 0.028583, -0.599263, -0.073634}},
       {{2.428247, 0.02662, -0.602419, -0.074123},
        {2.432657, 0.029055, -0.599293, -0.072732}},
       {{2.431515, 0.027687, -0.600096, -0.074459},
        {2.433075, 0.028543, -0.598987, -0.073985}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // multiple layers no norm
  model = TransformerDecoder(TransformerDecoderOptions(decoder_layer, 2));
  if (is_cuda) {
    model->to(torch::kCUDA);
  }

  decoder_input = torch::tensor({{{20, 30, 40, 50}}}, tensor_options);
  memory_input = torch::tensor({{{60, 70, 80, 90}}}, tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.31316, 0.0950293, -0.671995, 0.102802}}}, tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // multiple layers no norm
  model = TransformerDecoder(TransformerDecoderOptions(decoder_layer, 6));
  if (is_cuda) {
    model->to(torch::kCUDA);
  }
  // deterministic input
  decoder_input = torch::tensor(
      {{{0.4517, 0.6793, 0.5313, 0.0034}, {0.2678, 0.3677, 0.4459, 0.7166}},
       {{0.8100, 0.3716, 0.4096, 0.1976}, {0.6958, 0.8844, 0.6081, 0.8315}},
       {{0.0494, 0.9343, 0.5955, 0.3830}, {0.5404, 0.3464, 0.9378, 0.6200}}},
      tensor_options);
  memory_input = torch::tensor(
      {{{0.7462, 0.6653, 0.5679, 0.4891}, {0.5387, 0.1655, 0.3565, 0.0471}},
       {{0.8335, 0.2799, 0.5031, 0.2947}, {0.1402, 0.0318, 0.7636, 0.1346}},
       {{0.6333, 0.9344, 0.1376, 0.9938}, {0.8924, 0.2872, 0.6692, 0.2944}},
       {{0.9897, 0.6915, 0.3154, 0.1733}, {0.8645, 0.3513, 0.3064, 0.0767}},
       {{0.8117, 0.2366, 0.4838, 0.7881}, {0.3718, 0.4945, 0.9511, 0.0864}}},
      tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.42794, 0.026164, -0.60263, -0.0747591},
        {2.43113, 0.0279516, -0.600376, -0.0736896}},
       {{2.42794, 0.026164, -0.60263, -0.0747591},
        {2.43113, 0.0279516, -0.600376, -0.0736896}},
       {{2.42794, 0.026164, -0.60263, -0.0747591},
        {2.43113, 0.0279516, -0.600376, -0.0736896}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // multiple layers with norm
  LayerNorm norm(LayerNormOptions({decoder_layer.get()->options.d_model()}));
  model = TransformerDecoder(
      TransformerDecoderOptions(decoder_layer, 2).norm(AnyModule(norm)));
  if (is_cuda) {
    model->to(torch::kCUDA);
  }

  decoder_input = torch::tensor({{{20, 30, 40, 50}}}, tensor_options);
  memory_input = torch::tensor({{{60, 70, 80, 90}}}, tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{1.66166, -0.326986, -1.01466, -0.320017}}}, tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // multiple layers with norm
  model = TransformerDecoder(
      TransformerDecoderOptions(decoder_layer, 6).norm(AnyModule(norm)));
  if (is_cuda) {
    model->to(torch::kCUDA);
  }
  // deterministic input
  decoder_input = torch::tensor(
      {{{0.4517, 0.6793, 0.5313, 0.0034}, {0.2678, 0.3677, 0.4459, 0.7166}},
       {{0.8100, 0.3716, 0.4096, 0.1976}, {0.6958, 0.8844, 0.6081, 0.8315}},
       {{0.0494, 0.9343, 0.5955, 0.3830}, {0.5404, 0.3464, 0.9378, 0.6200}}},
      tensor_options);
  memory_input = torch::tensor(
      {{{0.7462, 0.6653, 0.5679, 0.4891}, {0.5387, 0.1655, 0.3565, 0.0471}},
       {{0.8335, 0.2799, 0.5031, 0.2947}, {0.1402, 0.0318, 0.7636, 0.1346}},
       {{0.6333, 0.9344, 0.1376, 0.9938}, {0.8924, 0.2872, 0.6692, 0.2944}},
       {{0.9897, 0.6915, 0.3154, 0.1733}, {0.8645, 0.3513, 0.3064, 0.0767}},
       {{0.8117, 0.2366, 0.4838, 0.7881}, {0.3718, 0.4945, 0.9511, 0.0864}}},
      tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{1.69559, -0.357291, -0.894741, -0.443553},
        {1.69571, -0.357363, -0.894154, -0.444196}},
       {{1.69559, -0.357291, -0.894741, -0.443553},
        {1.69571, -0.357363, -0.894154, -0.444196}},
       {{1.69559, -0.357291, -0.894741, -0.443553},
        {1.69571, -0.357363, -0.894154, -0.444196}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // gelu activation test cases
  decoder_layer.get()->options.activation(torch::kGELU);
  model = TransformerDecoder(TransformerDecoderOptions(decoder_layer, 1));
  if (is_cuda) {
    model->to(torch::kCUDA);
  }

  // deterministic input
  decoder_input = torch::tensor({{{20, 30, 40, 50}}}, tensor_options);
  memory_input = torch::tensor({{{60, 70, 80, 90}}}, tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.306435, 0.095946, -0.675796, 0.10687}}}, tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // deterministic input
  decoder_input =
      torch::tensor({{{9, 10, 11, 12}}, {{11, 12, 13, 14}}}, tensor_options);
  memory_input = torch::tensor({{{1, 2, 3, 4}}}, tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.415448, 0.054389, -0.610932, -0.0156613}},
       {{2.415448, 0.054389, -0.610932, -0.0156613}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // deterministic input
  decoder_input =
      torch::tensor({{{1, 2, 3, 4}}, {{5, 6, 7, 8}}}, tensor_options);
  memory_input =
      torch::tensor({{{9, 10, 11, 12}}, {{11, 12, 13, 14}}}, tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.338531, 0.087709, -0.65776, 0.080646}},
       {{2.338531, 0.087709, -0.65776, 0.080646}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // deterministic input
  decoder_input = torch::tensor(
      {{{0.4517, 0.6793, 0.5313, 0.0034}, {0.2678, 0.3677, 0.4459, 0.7166}},
       {{0.8100, 0.3716, 0.4096, 0.1976}, {0.6958, 0.8844, 0.6081, 0.8315}},
       {{0.0494, 0.9343, 0.5955, 0.3830}, {0.5404, 0.3464, 0.9378, 0.6200}}},
      tensor_options);
  memory_input = torch::tensor(
      {{{0.7462, 0.6653, 0.5679, 0.4891}, {0.5387, 0.1655, 0.3565, 0.0471}},
       {{0.8335, 0.2799, 0.5031, 0.2947}, {0.1402, 0.0318, 0.7636, 0.1346}},
       {{0.6333, 0.9344, 0.1376, 0.9938}, {0.8924, 0.2872, 0.6692, 0.2944}},
       {{0.9897, 0.6915, 0.3154, 0.1733}, {0.8645, 0.3513, 0.3064, 0.0767}},
       {{0.8117, 0.2366, 0.4838, 0.7881}, {0.3718, 0.4945, 0.9511, 0.0864}}},
      tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.42049104, 0.03443088, -0.60793706, -0.05436271},
        {2.42210631, 0.03546578, -0.60679895, -0.05357488}},
       {{2.41907674, 0.0336104, -0.60892977, -0.05490462},
        {2.42216881, 0.03586554, -0.6067524, -0.05289126}},
       {{2.42205716, 0.03488046, -0.60683681, -0.05460596},
        {2.42240309, 0.0354595, -0.60659063, -0.05378816}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // Multiple layers no norm
  model = TransformerDecoder(TransformerDecoderOptions(decoder_layer, 6));
  if (is_cuda) {
    model->to(torch::kCUDA);
  }
  decoder_input = torch::tensor(
      {{{0.4517, 0.6793, 0.5313, 0.0034}, {0.2678, 0.3677, 0.4459, 0.7166}},
       {{0.8100, 0.3716, 0.4096, 0.1976}, {0.6958, 0.8844, 0.6081, 0.8315}},
       {{0.0494, 0.9343, 0.5955, 0.3830}, {0.5404, 0.3464, 0.9378, 0.6200}}},
      tensor_options);
  memory_input = torch::tensor(
      {{{0.7462, 0.6653, 0.5679, 0.4891}, {0.5387, 0.1655, 0.3565, 0.0471}},
       {{0.8335, 0.2799, 0.5031, 0.2947}, {0.1402, 0.0318, 0.7636, 0.1346}},
       {{0.6333, 0.9344, 0.1376, 0.9938}, {0.8924, 0.2872, 0.6692, 0.2944}},
       {{0.9897, 0.6915, 0.3154, 0.1733}, {0.8645, 0.3513, 0.3064, 0.0767}},
       {{0.8117, 0.2366, 0.4838, 0.7881}, {0.3718, 0.4945, 0.9511, 0.0864}}},
      tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{2.41859, 0.0328114, -0.609269, -0.0560386},
        {2.42138, 0.034598, -0.607316, -0.0546574}},
       {{2.41859, 0.0328114, -0.609269, -0.0560386},
        {2.42138, 0.034598, -0.607316, -0.0546574}},
       {{2.41859, 0.0328114, -0.609269, -0.0560386},
        {2.42138, 0.034598, -0.607316, -0.0546574}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));

  // Multiple layers with norm
  norm = LayerNorm(LayerNormOptions({decoder_layer.get()->options.d_model()}));
  model = TransformerDecoder(
      TransformerDecoderOptions(decoder_layer, 6).norm(AnyModule(norm)));
  if (is_cuda) {
    model->to(torch::kCUDA);
  }

  decoder_input = torch::tensor(
      {{{0.4517, 0.6793, 0.5313, 0.0034}, {0.2678, 0.3677, 0.4459, 0.7166}},
       {{0.8100, 0.3716, 0.4096, 0.1976}, {0.6958, 0.8844, 0.6081, 0.8315}},
       {{0.0494, 0.9343, 0.5955, 0.3830}, {0.5404, 0.3464, 0.9378, 0.6200}}},
      tensor_options);
  memory_input = torch::tensor(
      {{{0.7462, 0.6653, 0.5679, 0.4891}, {0.5387, 0.1655, 0.3565, 0.0471}},
       {{0.8335, 0.2799, 0.5031, 0.2947}, {0.1402, 0.0318, 0.7636, 0.1346}},
       {{0.6333, 0.9344, 0.1376, 0.9938}, {0.8924, 0.2872, 0.6692, 0.2944}},
       {{0.9897, 0.6915, 0.3154, 0.1733}, {0.8645, 0.3513, 0.3064, 0.0767}},
       {{0.8117, 0.2366, 0.4838, 0.7881}, {0.3718, 0.4945, 0.9511, 0.0864}}},
      tensor_options);
  result = model(decoder_input, memory_input).detach();
  ref_output = torch::tensor(
      {{{1.69298, -0.355163, -0.906375, -0.431439},
        {1.69305, -0.355195, -0.906062, -0.431791}},
       {{1.69298, -0.355163, -0.906375, -0.431439},
        {1.69305, -0.355195, -0.906062, -0.431791}},
       {{1.69298, -0.355163, -0.906375, -0.431439},
        {1.69305, -0.355195, -0.906062, -0.431791}}},
      tensor_options);
  ASSERT_EQ(result.sizes().size(), ref_output.sizes().size());
  ASSERT_TRUE(torch::allclose(
      result,
      ref_output,
      1e-7,
      1e-5,
      /*equal_nan=*/true));
}

TEST_F(TransformerTest, TransformerDecoder) {
  transformer_decoder_test_helper(
      /*is_cuda=*/false, /*use_callable_activation=*/false);
  transformer_decoder_test_helper(
      /*is_cuda=*/false, /*use_callable_activation=*/true);
}

TEST_F(TransformerTest, TransformerDecoder_CUDA) {
  transformer_decoder_test_helper(
      /*is_cuda=*/true, /*use_callable_activation=*/false);
  transformer_decoder_test_helper(
      /*is_cuda=*/true, /*use_callable_activation=*/true);
}

TEST_F(TransformerTest, PrettyPrintTransformerDecoder) {
  LayerNorm norm = LayerNorm(LayerNormOptions({4}));
  TransformerDecoderOptions options(
      TransformerDecoderOptions(TransformerDecoderLayerOptions(4, 2), 2)
          .norm(AnyModule(norm)));
  ASSERT_EQ(
      c10::str(TransformerDecoder(options)),
      "torch::nn::TransformerDecoderImpl(\n"
      "  (layers): torch::nn::ModuleList(\n"
      "    (0): torch::nn::TransformerDecoderLayerImpl(\n"
      "      (self_attn): torch::nn::MultiheadAttention(\n"
      "        (out_proj): torch::nn::Linear(in_features=4, out_features=4, bias=true)\n"
      "      )\n"
      "      (multihead_attn): torch::nn::MultiheadAttention(\n"
      "        (out_proj): torch::nn::Linear(in_features=4, out_features=4, bias=true)\n"
      "      )\n"
      "      (linear1): torch::nn::Linear(in_features=4, out_features=2048, bias=true)\n"
      "      (dropout): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "      (linear2): torch::nn::Linear(in_features=2048, out_features=4, bias=true)\n"
      "      (norm1): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "      (norm2): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "      (norm3): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "      (dropout1): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "      (dropout2): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "      (dropout3): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "    )\n"
      "    (1): torch::nn::TransformerDecoderLayerImpl(\n"
      "      (self_attn): torch::nn::MultiheadAttention(\n"
      "        (out_proj): torch::nn::Linear(in_features=4, out_features=4, bias=true)\n"
      "      )\n"
      "      (multihead_attn): torch::nn::MultiheadAttention(\n"
      "        (out_proj): torch::nn::Linear(in_features=4, out_features=4, bias=true)\n"
      "      )\n"
      "      (linear1): torch::nn::Linear(in_features=4, out_features=2048, bias=true)\n"
      "      (dropout): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "      (linear2): torch::nn::Linear(in_features=2048, out_features=4, bias=true)\n"
      "      (norm1): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "      (norm2): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "      (norm3): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      "      (dropout1): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "      (dropout2): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "      (dropout3): torch::nn::Dropout(p=0.1, inplace=false)\n"
      "    )\n"
      "  )\n"
      "  (norm): torch::nn::LayerNorm([4], eps=1e-05, elementwise_affine=true)\n"
      ")");
}

void transformer_test_helper(bool is_cuda, bool use_callable_activation) {
  // this is a deterministic test for Transformere
  torch::Device device = is_cuda ? torch::kCUDA : torch::kCPU;
  torch::TensorOptions tensor_options =
      torch::TensorOptions().dtype(torch::kFloat32).device(device);

  // transformer created encoder/decoder
  auto options = TransformerOptions()
                     .d_model(4)
                     .nhead(2)
                     .num_encoder_layers(2)
                     .num_decoder_layers(1)
                     .dim_feedforward(16)
                     .dropout(0.0)
                     .activation(torch::kReLU);
  if (use_callable_activation) {
    options.activation(
        [&](const torch::Tensor& t) { return torch::nn::functional::relu(t); });
  }
  Transformer model(options);

  set_parameter_to_constants<Transformer>(model, tensor_options);
  if (tensor_options.device() == torch::kCUDA) {
    model->to(torch::kCUDA);
  }

  // transformer with customized encoder/decoder
  LayerNorm enorm(LayerNormOptions({4}));
  TransformerEncoder encoder(
      TransformerEncoderOptions(
          TransformerEncoderLayerOptions(4, 2).dim_feedforward(16).dropout(0.0),
          2)
          .norm(AnyModule(enorm)));

  LayerNorm dnorm(LayerNormOptions({4}));
  TransformerDecoder decoder(
      TransformerDecoderOptions(
          TransformerDecoderLayerOptions(4, 2).dim_feedforward(16).dropout(0.0),
          1)
          .norm(AnyModule(dnorm)));

  Transformer model_cus(TransformerOptions()
                            .d_model(4)
                            .nhead(2)
                            .custom_encoder(AnyModule(encoder))
                            .custom_decoder(AnyModule(decoder)));

  set_parameter_to_constants<Transformer>(model_cus, tensor_options);
  if (tensor_options.device() == torch::kCUDA) {
    model_cus->to(torch::kCUDA);
  }

  // test cases
  torch::Tensor src = torch::tensor(
      {{{1.0, 2.0, 3.0, 4.0}, {5.0, 6.0, 7.0, 8.0}},
       {{9.0, 10.0, 11.0, 12.0}, {13.0, 14.0, 15.0, 16.0}},
       {{17.0, 18.0, 19.0, 20.0}, {21.0, 22.0, 23.0, 24.0}}},
      tensor_options);

  torch::Tensor tgt = torch::tensor(
      {{{1.0, 2.0, 3.0, 4.0}, {5.0, 6.0, 7.0, 8.0}},
       {{9.0, 10.0, 11.0, 12.0}, {13.0, 14.0, 15.0, 16.0}}},
      tensor_options);

  torch::Tensor ref_output = torch::tensor(
      {{{2.695875, 0.347114, -0.044355, -0.549541},
        {2.696091, 0.347015, -0.044770, -0.548522}},
       {{2.695875, 0.347114, -0.044355, -0.549541},
        {2.696091, 0.347015, -0.044770, -0.548522}}},
      tensor_options);
  torch::Tensor result = model(src, tgt);
  torch::Tensor result_cus = model_cus(src, tgt);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(result.equal(result_cus));
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  torch::Tensor src_mask =
      Transformer::Impl::generate_square_subsequent_mask(src.size(0))
          .to(tensor_options);
  ref_output = torch::tensor(
      {{{2.695875, 0.347114, -0.044355, -0.549541},
        {2.696091, 0.347015, -0.044770, -0.548522}},
       {{2.695875, 0.347114, -0.044355, -0.549541},
        {2.696091, 0.347015, -0.044770, -0.548522}}},
      tensor_options);
  result = model(src, tgt, src_mask);
  result_cus = model_cus(src, tgt, src_mask);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(result.equal(result_cus));
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));

  torch::Tensor tgt_key_padding_mask =
      torch::zeros({tgt.size(1), tgt.size(0)}, tensor_options) == 1;
  tgt_key_padding_mask[0][0] = 1;
  tgt_key_padding_mask[1][1] = 1;
  ref_output = torch::tensor(
      {{{2.696114, 0.347004, -0.044813, -0.548417},
        {2.696091, 0.347015, -0.044770, -0.548522}},
       {{2.696114, 0.347004, -0.044813, -0.548417},
        {2.696091, 0.347015, -0.044770, -0.548522}}},
      tensor_options);
  result = model(
      src,
      tgt,
      src_mask,
      torch::Tensor(),
      torch::Tensor(),
      torch::Tensor(),
      tgt_key_padding_mask);
  result_cus = model_cus(
      src,
      tgt,
      src_mask,
      torch::Tensor(),
      torch::Tensor(),
      torch::Tensor(),
      tgt_key_padding_mask);
  ASSERT_EQ(result.sizes(), ref_output.sizes());
  ASSERT_TRUE(result.equal(result_cus));
  ASSERT_TRUE(
      torch::allclose(result, ref_output, 1e-7, 1e-5, /*equal_nan=*/true));
}

TEST_F(TransformerTest, Transformer) {
  transformer_test_helper(/*is_cuda=*/false, /*use_callable_activation=*/false);
  transformer_test_helper(/*is_cuda=*/false, /*use_callable_activation=*/true);
}

TEST_F(TransformerTest, Transformer_CUDA) {
  transformer_test_helper(/*is_cuda=*/true, /*use_callable_activation=*/false);
  transformer_test_helper(/*is_cuda=*/true, /*use_callable_activation=*/true);
}

TEST_F(TransformerTest, TransformerArgsCorrectness) {
  Transformer model(TransformerOptions()
                        .d_model(4)
                        .nhead(2)
                        .num_encoder_layers(2)
                        .num_decoder_layers(1)
                        .dim_feedforward(16)
                        .dropout(0.0)
                        .activation(torch::kReLU));

  torch::Tensor src = torch::randn({2, 3, 4});
  torch::Tensor tgt = torch::randn({3, 2, 4});

  ASSERT_THROWS_WITH(
      model(src, tgt), "src and tgt should have equal batch size");

  tgt = torch::randn({2, 3, 3});
  ASSERT_THROWS_WITH(
      model(src, tgt), "src and tgt should have same feature size as d_model");

  src = torch::randn({2, 3});
  ASSERT_THROWS_WITH(model(src, tgt), "src and tgt should have 3 dimensions");
}