Update

  output_ar = torch.tensor([[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, 0.0,

                 adj_matrix, distance_matrix)

  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)

  Examples

  --------

  >>> import deepchem as dc

<<<<<<< HEAD

  >>> block = dc.models.torch_models.layers.MATEncoder(dist_kernel = 'softmax', lambda_attention = 0.33, lambda_adistance = 0.33, h = 8, sa_hsize = 1024, sa_dropout_p = 0.1, d_input = 1024, activation = 'relu', n_layers = 1, ff_dropout_p = 0.1, encoder_hsize = 1024, encoder_dropout_p = 0.1, N = 3)

=======

  >>> from rdkit import Chem

  >>> mol = rdkit.Chem.MolFromSmiles("CC")

  >>> adj_matrix = Chem.GetAdjacencyMatrix(mol)

  >>> distance_matrix = Chem.GetDistanceMatrix(mol)

  >>> layer = dc.models.torch_models.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, 0.0, adj_matrix, distance_matrix)

>>>>>>> Update

  """

<<<<<<< HEAD