def get_mask(valid_length, full_length):
  """Returns a mask given the valid length of a vector with trailing padding.

  This is useful for masking out padded elements from a loss. For example
  ```
  input_graphs_tuple = ...
  input_graphs_tuple_size = get_graphs_tuple_size(input_graphs_tuple)
  padded_input_graphs_tuple = pad_graphs_tuple(input_graphs_tuple,
       pad_nodes_to,...)


  per_node_loss # After graph_network computation.

  nodes_mask = get_mask(
      input_graphs_tuple_size.nodes, pad_nodes_to)

  masked_per_node_loss = per_node_loss * tf.cast(
      nodes_mask, per_node_loss.dtype)
  ```

  Args:
    valid_length: Length of the valid elements.
    full_length: Full length of the vector after padding.

  Returns:
    Boolean mask of shape [full_length], where all values are set to `True`
    except for the last `padding_length` which are set to False.
  """
  valid_length = tf.cast(valid_length, tf.int32)
  full_length = tf.cast(full_length, tf.int32)
  # This implementation allows for statically sized shapes, rather than
  # using concat([tf.ones([valid_length]), tf.zeros([full_length-valid_length])]
  # which has intermediate tensors with shapes not know statically.
  field_mask = tf.range(full_length)
  field_mask = field_mask < valid_length
  return field_mask
Returns a mask given the valid length of a vector with trailing padding.

  This is useful for masking out padded elements from a loss. For example
  ```
  input_graphs_tuple = ...
  input_graphs_tuple_size = get_graphs_tuple_size(input_graphs_tuple)
  padded_input_graphs_tuple = pad_graphs_tuple(input_graphs_tuple,
       pad_nodes_to,...)


  per_node_loss # After graph_network computation.

  nodes_mask = get_mask(
      input_graphs_tuple_size.nodes, pad_nodes_to)

  masked_per_node_loss = per_node_loss * tf.cast(
      nodes_mask, per_node_loss.dtype)
  ```

  Args:
    valid_length: Length of the valid elements.
    full_length: Full length of the vector after padding.

  Returns:
    Boolean mask of shape [full_length], where all values are set to `True`
    except for the last `padding_length` which are set to False.
  
graph_nets/utils_tf.py