Merge pull request #2601 from atreyamaj/firstlayer

except:

  has_tensorflow = False

try:

  import torch

  import deepchem.models.torch_models.layers as torch_layers

  has_torch = True

except:

  has_torch = False

@pytest.mark.tensorflow

def test_cosine_dist():

  atom_features = np.random.rand(batch_size, n_atom_feat)

  membership = np.sort(np.random.randint(0, batch_size, size=(batch_size)))

  outputs = layer([atom_features, membership])

@pytest.mark.pytorch

def test_layer_norm():

  """Test invoking LayerNorm."""

  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)

import math

try:

  import torch

  import torch.nn as nn

except:

  raise ImportError('These classes require Torch to be installed.')

class ScaleNorm(nn.Module):

  """Apply Scale Normalization to input.

  The ScaleNorm layer first computes the square root of the scale, then computes the matrix/vector norm of the input tensor.

  The norm value is calculated as `sqrt(scale) / matrix norm`.

  Finally, the result is returned as `input_tensor * norm value`.

  This layer can be used instead of LayerNorm when a scaled version of the norm is required.

  Instead of performing the scaling operation (scale / norm) in a lambda-like layer, we are defining it within this layer to make prototyping more efficient.

  References

  ----------

  .. [1] Lukasz Maziarka et al. "Molecule Attention Transformer" Graph Representation Learning workshop and Machine Learning and the Physical Sciences workshop at NeurIPS 2019. 2020. https://arxiv.org/abs/2002.08264

  Examples

  --------

  >>> import deepchem as dc

  >>> scale = 0.35

  >>> layer = dc.models.torch_models.ScaleNorm(scale)

  >>> input_tensor = torch.tensor([[1.269, 39.36], [0.00918, -9.12]])

  >>> output_tensor = layer.forward(input_tensor)

  """

  def __init__(self, scale, eps=1e-5):

    """Initialize a ScaleNorm layer.

    Parameters

    ----------

    scale: Real number or single element tensor

      Scale magnitude.

    eps: float

      Epsilon value.

    """

    super(ScaleNorm, self).__init__()

    self.scale = nn.Parameter(torch.tensor(math.sqrt(scale)))

    self.eps = eps

  def forward(self, x):

    norm = self.scale / torch.norm(x, dim=-1, keepdim=True).clamp(min=self.eps)

    return x * norm

.. autoclass:: deepchem.models.layers.SetGather

  :members:

.. autoclass:: deepchem.models.torch_models.layers.ScaleNorm

  :members:

.. autofunction:: deepchem.models.layers.cosine_dist