def specs_from_graphs_tuple(
    graphs_tuple_sample,
    dynamic_num_graphs=False,
    dynamic_num_nodes=True,
    dynamic_num_edges=True,
    description_fn=tf.TensorSpec,
    ):
  """Returns the `TensorSpec` specification for a given `GraphsTuple`.

  This method is often used with `tf.function` in Tensorflow 2 to obtain
  improved speed and performance of eager code. For example:

  ```
  example_graphs_tuple = get_graphs_tuple(...)

  @tf.function(input_signature=[specs_from_graphs_tuple(example_graphs_tuple)])
  def forward_pass(graphs_tuple_input):
    graphs_tuple_output = graph_network(graphs_tuple_input)
    return graphs_tuple_output

  for i in range(num_training_steps):
    input = get_graphs_tuple(...)
    with tf.GradientTape() as tape:
      output = forward_pass(input)
      loss = compute_loss(output)
    grads = tape.gradient(loss, graph_network.trainable_variables)
    optimizer.apply(grads, graph_network.trainable_variables)
  ```

  Args:
    graphs_tuple_sample: A `graphs.GraphsTuple` with sample data. `GraphsTuple`s
      that have fields with `None` are not accepted since they will create an
      invalid signature specification for `tf.function`. If your graph has
      `None`s use `utils_tf.set_zero_edge_features`,
      `utils_tf.set_zero_node_features` or `utils_tf.set_zero_global_features`.
      This method also returns the signature for `GraphTuple`s with nests of
      tensors in the feature fields (`nodes`, `edges`, `globals`), including
      empty nests (e.g. empty list, dict, or tuple). Nested node, edge and
      global feature tensors, should usually have the same leading dimension as
      all other node, edge and global feature tensors respectively.
    dynamic_num_graphs: Boolean indicating if the number of graphs in each
      `GraphsTuple` will be variable across examples.
    dynamic_num_nodes: Boolean indicating if number of nodes per graph will be
      variable across examples. Not used if `dynamic_num_graphs` is True, as the
      size of the first axis of all `GraphsTuple` fields will be variable, due
      to the variable number of graphs.
    dynamic_num_edges: Boolean indicating if number of edges per graph will be
      variable across examples. Not used if dynamic_num_graphs is True, as the
      size of the first axis of all `GraphsTuple` fields will be variable, due
      to the variable number of graphs.
    description_fn: A callable which accepts the dtype and shape arguments to
      describe the shapes and types of tensors. By default uses `tf.TensorSpec`.

  Returns:
    A `GraphsTuple` with tensors replaced by `TensorSpec` with shape and dtype
    of the field contents.

  Raises:
    ValueError: If a `GraphsTuple` has a field with `None`.
  """
  graphs_tuple_description_fields = {}
  edge_dim_fields = [graphs.EDGES, graphs.SENDERS, graphs.RECEIVERS]

  # Method to get the spec for a single tensor.
  def get_tensor_spec(tensor, field_name):
    """Returns the spec of an array or a tensor in the field of a graph."""

    shape = list(tensor.shape)
    dtype = tensor.dtype

    # If the field is not None but has no field shape (i.e. it is a constant)
    # then we consider this to be a replaced `None`.
    # If dynamic_num_graphs, then all fields have a None first dimension.
    # If dynamic_num_nodes, then the "nodes" field needs None first dimension.
    # If dynamic_num_edges, then the "edges", "senders" and "receivers" need
    # a None first dimension.
    if (shape and (
        dynamic_num_graphs or
        (dynamic_num_nodes and field_name == graphs.NODES) or
        (dynamic_num_edges and field_name in edge_dim_fields))):
      shape[0] = None
    return description_fn(shape=shape, dtype=dtype)

  for field_name in graphs.ALL_FIELDS:
    field_sample = getattr(graphs_tuple_sample, field_name)
    if field_sample is None:
      raise ValueError(
          "The `GraphsTuple` field `{}` was `None`. All fields of the "
          "`GraphsTuple` must be specified to create valid signatures that"
          "work with `tf.function`. This can be achieved with `input_graph = "
          "utils_tf.set_zero_{{node,edge,global}}_features(input_graph, 0)`"
          "to replace None's by empty features in your graph. Alternatively"
          "`None`s can be replaced by empty lists by doing `input_graph = "
          "input_graph.replace({{nodes,edges,globals}}=[]). To ensure "
          "correct execution of the program, it is recommended to restore "
          "the None's once inside of the `tf.function` by doing "
          "`input_graph = input_graph.replace({{nodes,edges,globals}}=None)"
          "".format(field_name))

    if field_name in graphs.GRAPH_FEATURE_FIELDS:
      field_spec = tree.map_structure(
          functools.partial(get_tensor_spec, field_name=field_name),
          field_sample)
    else:
      field_spec = get_tensor_spec(field_sample, field_name=field_name)

    graphs_tuple_description_fields[field_name] = field_spec

  return graphs.GraphsTuple(**graphs_tuple_description_fields)
Returns the `TensorSpec` specification for a given `GraphsTuple`.

  This method is often used with `tf.function` in Tensorflow 2 to obtain
  improved speed and performance of eager code. For example:

  ```
  example_graphs_tuple = get_graphs_tuple(...)

  @tf.function(input_signature=[specs_from_graphs_tuple(example_graphs_tuple)])
  def forward_pass(graphs_tuple_input):
    graphs_tuple_output = graph_network(graphs_tuple_input)
    return graphs_tuple_output

  for i in range(num_training_steps):
    input = get_graphs_tuple(...)
    with tf.GradientTape() as tape:
      output = forward_pass(input)
      loss = compute_loss(output)
    grads = tape.gradient(loss, graph_network.trainable_variables)
    optimizer.apply(grads, graph_network.trainable_variables)
  ```

  Args:
    graphs_tuple_sample: A `graphs.GraphsTuple` with sample data. `GraphsTuple`s
      that have fields with `None` are not accepted since they will create an
      invalid signature specification for `tf.function`. If your graph has
      `None`s use `utils_tf.set_zero_edge_features`,
      `utils_tf.set_zero_node_features` or `utils_tf.set_zero_global_features`.
      This method also returns the signature for `GraphTuple`s with nests of
      tensors in the feature fields (`nodes`, `edges`, `globals`), including
      empty nests (e.g. empty list, dict, or tuple). Nested node, edge and
      global feature tensors, should usually have the same leading dimension as
      all other node, edge and global feature tensors respectively.
    dynamic_num_graphs: Boolean indicating if the number of graphs in each
      `GraphsTuple` will be variable across examples.
    dynamic_num_nodes: Boolean indicating if number of nodes per graph will be
      variable across examples. Not used if `dynamic_num_graphs` is True, as the
      size of the first axis of all `GraphsTuple` fields will be variable, due
      to the variable number of graphs.
    dynamic_num_edges: Boolean indicating if number of edges per graph will be
      variable across examples. Not used if dynamic_num_graphs is True, as the
      size of the first axis of all `GraphsTuple` fields will be variable, due
      to the variable number of graphs.
    description_fn: A callable which accepts the dtype and shape arguments to
      describe the shapes and types of tensors. By default uses `tf.TensorSpec`.

  Returns:
    A `GraphsTuple` with tensors replaced by `TensorSpec` with shape and dtype
    of the field contents.

  Raises:
    ValueError: If a `GraphsTuple` has a field with `None`.
  
graph_nets/utils_tf.py