adding tests

  """Test invoking ScaleNorm."""

  input_ar = torch.tensor([[1., 99., 10000.], [0.003, 999.37, 23.]])

  layer = torch_layers.ScaleNorm(0.35)

  result1 = layer(input_ar)

  output_ar = torch.tensor([[5.9157897e-05, 5.8566318e-03, 5.9157896e-01],

  result1 = layer.forward(input_ar)

  output_ar = np.array([[5.9157897e-05, 5.8566318e-03, 5.9157896e-01],

                        [1.7754727e-06, 5.9145141e-01, 1.3611957e-02]])

  assert torch.allclose(result1, output_ar)

@pytest.mark.torch

def test_multi_headed_mat_attention():

  """Test invoking MultiHeadedMATAttention."""

  from rdkit import Chem

  torch.manual_seed(0)

  input_smile = "CC"

  mol = Chem.MolFromSmiles(input_smile)

  adj_matrix = Chem.GetAdjacencyMatrix(mol)

  distance_matrix = Chem.GetDistanceMatrix(mol)

  layer = torch_layers.MultiHeadedMATAttention(

      dist_kernel='softmax',

      lambda_attention=0.33,

      lambda_distance=0.33,

      h=2,

      hsize=2,

      dropout_p=0.0)

  input_tensor = torch.tensor([[1., 2.], [5., 6.]])

  mask = torch.tensor([[1., 1.], [1., 1.]])

  result = layer(input_tensor, input_tensor, input_tensor, mask, adj_matrix,

                 distance_matrix, 0.0)

  output_ar = torch.tensor([[[0.0492, -0.0792], [-0.9971, -0.3172],

                             [0.0492, -0.0792], [-0.9971, -0.3172]],

                            [[0.8671, 0.1069], [-3.4075, -0.8656],

                             [0.8671, 0.1069], [-3.4075, -0.8656]]])

  assert torch.allclose(result, output_ar, rtol=1e-3)

  assert np.allclose(result1, output_ar)

@pytest.mark.torch

def test_position_wise_feed_forward():

  """Test invoking PositionwiseFeedForward."""

def test_mat_embedding():

  """Test invoking MATEmbedding."""

  torch.manual_seed(0)

  input_ar = torch.tensor([[1., 2.], [5., 6.]])

  layer = torch_layers.PositionwiseFeedForward(

      d_input=2,

      d_hidden=2,

      d_output=2,

      activation='relu',

      n_layers=1,

      dropout_p=0.0)

  result = layer(input_ar)

  output_ar = torch.tensor([[0.4810, 0.0000], [1.9771, 0.0000]])

  assert torch.allclose(result, output_ar, rtol=1e-4)

@pytest.mark.torch

def test_sub_layer_connection():

  """Test invoking SublayerConnection."""

  torch.manual_seed(0)

  input_ar = torch.tensor([[1., 2.], [5., 6.]])

  layer = torch_layers.SublayerConnection(2, 0.0)

  result = layer(input_ar, input_ar)

  output_ar = torch.tensor([[2.0027e-05, 3.0000e+00], [4.0000e+00, 7.0000e+00]])

  assert torch.allclose(result, output_ar)

  input_ar = torch.tensor([1.,2.,3.])

  layer = torch_layers.MATEmbedding(3, 1, 0.0)

  result1 = layer(input_ar)

  output_ar = np.array([[5.9157897e-05, 5.8566318e-03, 5.9157896e-01],

                        [1.7754727e-06, 5.9145141e-01, 1.3611957e-02]])

  assert np.allclose(result1, output_ar)

@pytest.mark.torch

def test_mat_encoder_layer():

  """Test invoking MATEncoderLayer."""

  torch.manual_seed(0)

  from rdkit import Chem

  input_ar = torch.Tensor([[1., 2.], [5., 6.]])

  mask = torch.Tensor([[1., 1.], [1., 1.]])

  mol = Chem.MolFromSmiles("CC")

  adj_matrix = Chem.GetAdjacencyMatrix(mol)

  distance_matrix = Chem.GetDistanceMatrix(mol)

  layer = torch_layers.MATEncoderLayer(

      dist_kernel='softmax',

      lambda_attention=0.33,

      lambda_distance=0.33,

      h=2,

      sa_hsize=2,

      sa_dropout_p=0.0,

      output_bias=True,

      d_input=2,

      d_hidden=2,

      d_output=2,

      activation='relu',

      n_layers=2,

      ff_dropout_p=0.0,

      encoder_hsize=2,

      encoder_dropout_p=0.0)

  result = layer(input_ar, mask, adj_matrix, distance_matrix, 0.0)

  output_ar = torch.tensor([[[0.9988, 2.0012], [-0.9999, 3.9999],

                             [0.9988, 2.0012], [-0.9999, 3.9999]],

                            [[5.0000, 6.0000], [3.0000, 8.0000],

                             [5.0000, 6.0000], [3.0000, 8.0000]]])

  assert torch.allclose(result, output_ar, rtol=1e-4)

def test_mat_generator():

  """Test invoking MATGenerator."""

  input_ar = torch.tensor([[1., 99., 10000.], [0.003, 999.37, 23.]])

  layer = torch_layers.ScaleNorm(0.35)

  result1 = layer.forward(input_ar)

  output_ar = np.array([[5.9157897e-05, 5.8566318e-03, 5.9157896e-01],

                        [1.7754727e-06, 5.9145141e-01, 1.3611957e-02]])

  assert np.allclose(result1, output_ar)