Method map

graph_nets/graphs.py:178–195

Applies `field_fn` to the fields `fields` of the instance. `field_fn` is applied exactly once per field in `fields`. The result must satisfy the `GraphsTuple` requirement w.r.t. `None` fields, i.e. the `SENDERS` cannot be `None` if the `EDGES` or `RECEIVERS` are not `None`, etc.

(self, field_fn, fields=GRAPH_FEATURE_FIELDS)

    return output

  def map(self, field_fn, fields=GRAPH_FEATURE_FIELDS):
    """Applies `field_fn` to the fields `fields` of the instance.

    `field_fn` is applied exactly once per field in `fields`. The result must
    satisfy the `GraphsTuple` requirement w.r.t. `None` fields, i.e. the
    `SENDERS` cannot be `None` if the `EDGES` or `RECEIVERS` are not `None`,
    etc.

    Args:
      field_fn: A callable that take a single argument.
      fields: (iterable of `str`). An iterable of the fields to apply
        `field_fn` to.

    Returns:
      A copy of the instance, with the fields in `fields` replaced by the result
      of applying `field_fn` to them.
    """
    return self.replace(**{k: field_fn(getattr(self, k)) for k in fields})