added docstrings

    else:

      self.variable_scope = local_scope

  def set_summary(self, summary_op, summary_description=None, collections=None):

  def set_summary(self, summary_op, summary_description=None,

                  collections=None):

    """Annotates a tensor with a tf.summary operation

    Collects data from self.out_tensor by default but can be changed by setting 

    self.tb_input to another tensor in create_tensor

    elif isinstance(in_layer, tf.Tensor):

      layers.append(TensorWrapper(in_layer))

    else:

      raise ValueError("convert_to_layers must be invoked on layers or tensors")

      raise ValueError(

          "convert_to_layers must be invoked on layers or tensors")

  return layers

class Dense(Layer):

  def __init__(

      self,

      out_channels,

class Reshape(Layer):

  def __init__(self, shape, **kwargs):

    super(Reshape, self).__init__(**kwargs)

    self._new_shape = tuple(-1 if x is None else x for x in shape)

class Squeeze(Layer):

  def __init__(self, in_layers=None, squeeze_dims=None, **kwargs):

    self.squeeze_dims = squeeze_dims

    super(Squeeze, self).__init__(in_layers, **kwargs)

class Transpose(Layer):

  def __init__(self, perm, **kwargs):

    super(Transpose, self).__init__(**kwargs)

    self.perm = perm

class CombineMeanStd(Layer):

  def __init__(self, in_layers=None, **kwargs):

    super(CombineMeanStd, self).__init__(in_layers, **kwargs)

    try:

class Repeat(Layer):

  def __init__(self, n_times, **kwargs):

    self.n_times = n_times

    super(Repeat, self).__init__(**kwargs)

class TimeSeriesDense(Layer):

  def __init__(self, out_channels, **kwargs):

    self.out_channels = out_channels

    super(TimeSeriesDense, self).__init__(**kwargs)

class Input(Layer):

  def __init__(self, shape, dtype=tf.float32, **kwargs):

    self._shape = tuple(shape)

    self.dtype = dtype

class Feature(Input):

  def __init__(self, **kwargs):

    super(Feature, self).__init__(**kwargs)

class Label(Input):

  def __init__(self, **kwargs):

    super(Label, self).__init__(**kwargs)

class Weights(Input):

  def __init__(self, **kwargs):

    super(Weights, self).__init__(**kwargs)

class L1Loss(Layer):

  def __init__(self, in_layers=None, **kwargs):

    super(L1Loss, self).__init__(in_layers, **kwargs)

class L2Loss(Layer):

  def __init__(self, in_layers=None, **kwargs):

    super(L2Loss, self).__init__(in_layers, **kwargs)

    try:

class SoftMax(Layer):

  def __init__(self, in_layers=None, **kwargs):

    super(SoftMax, self).__init__(in_layers, **kwargs)

    try:

class Concat(Layer):

  def __init__(self, in_layers=None, axis=1, **kwargs):

    self.axis = axis

    super(Concat, self).__init__(in_layers, **kwargs)

class Stack(Layer):

  def __init__(self, in_layers=None, axis=1, **kwargs):

    self.axis = axis

    super(Stack, self).__init__(in_layers, **kwargs)

class SparseSoftMaxCrossEntropy(Layer):

  def __init__(self, in_layers=None, **kwargs):

    super(SparseSoftMaxCrossEntropy, self).__init__(in_layers, **kwargs)

    try:

class SoftMaxCrossEntropy(Layer):

  def __init__(self, in_layers=None, **kwargs):

    super(SoftMaxCrossEntropy, self).__init__(in_layers, **kwargs)

    try:

class ReduceMean(Layer):

  def __init__(self, in_layers=None, axis=None, **kwargs):

    self.axis = axis

    super(ReduceMean, self).__init__(in_layers, **kwargs)

class ToFloat(Layer):

  def __init__(self, in_layers=None, **kwargs):

    super(ToFloat, self).__init__(in_layers, **kwargs)

    try:

class ReduceSum(Layer):

  def __init__(self, in_layers=None, axis=None, **kwargs):

    self.axis = axis

    super(ReduceSum, self).__init__(in_layers, **kwargs)

class ReduceSquareDifference(Layer):

  def __init__(self, in_layers=None, axis=None, **kwargs):

    self.axis = axis

    super(ReduceSquareDifference, self).__init__(in_layers, **kwargs)

class MaxPool(Layer):

  def __init__(self,

               ksize=[1, 2, 2, 1],

               strides=[1, 2, 2, 1],

    inputs = self._get_input_tensors(in_layers)

    in_tensor = inputs[0]

    out_tensor = tf.nn.max_pool(

        in_tensor, ksize=self.ksize, strides=self.strides, padding=self.padding)

        in_tensor,

        ksize=self.ksize,

        strides=self.strides,

        padding=self.padding)

    if set_tensors:

      self.out_tensor = out_tensor

    return out_tensor

class GraphConv(Layer):

  def __init__(self,

               out_channel,

               min_deg=0,

class GraphPool(Layer):

  def __init__(self, min_degree=0, max_degree=10, **kwargs):

    self.min_degree = min_degree

    self.max_degree = max_degree

class GraphGather(Layer):

  def __init__(self, batch_size, activation_fn=None, **kwargs):

    self.batch_size = batch_size

    self.activation_fn = activation_fn

  y: tf.Tensor

    Input Tensor 

  """

  denom = (

      model_ops.sqrt(model_ops.sum(tf.square(x)) * model_ops.sum(tf.square(y)))

      + model_ops.epsilon())

  denom = (model_ops.sqrt(

      model_ops.sum(tf.square(x)) * model_ops.sum(tf.square(y))) +

           model_ops.epsilon())

  return model_ops.dot(x, tf.transpose(y)) / denom

class BatchNorm(Layer):

  def __init__(self, in_layers=None, **kwargs):

    super(BatchNorm, self).__init__(in_layers, **kwargs)

    try:

class BatchNormalization(Layer):

  def __init__(self,

               epsilon=1e-5,

               axis=-1,

class WeightedError(Layer):

  def __init__(self, in_layers=None, **kwargs):

    super(WeightedError, self).__init__(in_layers, **kwargs)

    self._shape = tuple()

    # List of length N_atoms, each element of different length uniques_i

    nbrs = self.get_atoms_in_nbrs(coords, cells)

    padding = tf.fill((self.M_nbrs, ), -1)

    padded_nbrs = [tf.concat([unique_nbrs, padding], 0) for unique_nbrs in nbrs]

    padded_nbrs = [

        tf.concat([unique_nbrs, padding], 0) for unique_nbrs in nbrs

    ]

    # List of length N_atoms, each element of different length uniques_i

    # List of length N_atoms, each a tensor of shape

class Dropout(Layer):

  def __init__(self, dropout_prob, **kwargs):

    self.dropout_prob = dropout_prob

    super(Dropout, self).__init__(**kwargs)

  def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):

    inputs = self._get_input_tensors(in_layers)

    parent_tensor = inputs[0]

    out_tensor = parent_tensor + model_ops.weight_decay(self.penalty_type,

                                                        self.penalty)

    out_tensor = parent_tensor + model_ops.weight_decay(

        self.penalty_type, self.penalty)

    if set_tensors:

      self.out_tensor = out_tensor

    return out_tensor

class AtomicConvolution(Layer):

  def __init__(self,

               atom_types=None,

               radial_params=list(),

    example_tensors = tf.unstack(X, axis=0)

    example_nbrs = tf.unstack(nbr_indices, axis=0)

    all_nbr_coords = []

    for example, (example_tensor,

    for example, (

        example_tensor,

        example_nbr) in enumerate(zip(example_tensors, example_nbrs)):

      nbr_coords = tf.gather(example_tensor, example_nbr)

      all_nbr_coords.append(nbr_coords)

    return R

def AlphaShare(**kwargs):

def AlphaShare(in_layers=None, **kwargs):

  """

  This method should be used when constructing AlphaShare layers from Sluice Networks

  Parameters

  ----------

  in_layers: list of Layers or tensors

    tensors in list must be the same size and list must include two or more tensors

  Returns

  -------

  output_layers: list of Layers or tensors with same size as in_layers

    Distance matrix.

  References:

  Sluice networks: Learning what to share between loosely related tasks

  https://arxiv.org/abs/1705.08142

  """

  output_layers = []

  alpha_share = AlphaShareLayer(**kwargs)

  alpha_share = AlphaShareLayer(in_layers=in_layers, **kwargs)

  num_outputs = len(kwargs['in_layers'])

  for num_layer in range(0, num_outputs):

    output_layers.append(

  """

  Part of a sluice network. Adds alpha parameters to control

  sharing between the main and auxillary tasks

  Factory method AlphaShare should be used for construction

  Parameters

  ----------

  in_layers: list of Layers or tensors

    tensors in list must be the same size and list must include two or more tensors

  Returns

  -------

  out_tensor: a tensor with shape [len(in_layers), x, y] where x, y were the original layer dimensions

    out_tensor should be fed into LayerSplitter 

  Distance matrix.

  """

  def __init__(self, **kwargs):

  """

  Part of a sluice network. Adds beta params to control which layer

  outputs are used for prediction

  Parameters

  ----------

  in_layers: list of Layers or tensors

    tensors in list must be the same size and list must include two or more tensors

  Returns

  -------

  output_layers: list of Layers or tensors with same size as in_layers

    Distance matrix.

  """

  def __init__(self, **kwargs):