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)
Returns the spec of an array or a tensor in the field of a graph.
graph_nets/utils_tf.py