compiler

Defines the mG compiler and the basic mG functions.

This package defines a compiler for mG programs and the data structures to instantiate it. It also provides the various classes that allow for the definition of mG functions.

The package contains the following classes:

• NodeConfig
• EdgeConfig
• CompilerConfig
• MGCompiler
• PsiNonLocal
• PsiLocal
• PsiGlobal
• Phi
• Sigma
• Constant
• Pi

The module contains the following functions:

• make_uoperator(op, name)
• make_boperator(op, name)
• make_koperator(op, name)

PsiLocal

PsiLocal(f: Callable | None = None, name: str | None = None)

Bases: Psi

A psi function of the mG language that only applies a local transformation of node labels.

A local transformation of node labels \(\psi: T \rightarrow U\).

Examples:

>>> PsiLocal(lambda x: x + 1, name='Add1')
<PsiLocal ...>

Parameters:

• f (Callable | None, default: None ) –

  The function that this object will run when called. The function must be compatible with TensorFlow's broadcasting rules. The function is expected to take in input a tensor X with shape (n_nodes, n_node_features), containing the labels of every node in the graph, and return a tensor of shape (n_nodes, n_new_node_features), containing the transformed labels. The function is not supposed to use global information, but this is not enforced.

• name (str | None, default: None ) –

  The name of the function.

Source code in libmg/compiler/functions.py

def __init__(self, f: Callable | None = None, name: str | None = None):
    """Initializes the instance with a function and a name.

    Args:
        f: The function that this object will run when called. The function must be compatible with TensorFlow's broadcasting
            rules. The function is expected to take in input a tensor X with shape ``(n_nodes, n_node_features)``, containing the labels of
            every node in the graph, and return a tensor of shape ``(n_nodes, n_new_node_features)``, containing the transformed labels.
            The function is not supposed to use global information, but this is not enforced.
        name: The name of the function.
    """
    if f is None:
        f = self.func
    super().__init__(f, name)

fname property

fname

The name of this function. This can be either the name provided during initialization, if it was provided, or the dynamic class name.

make classmethod

make(name: str | None, f: Callable) -> Callable[[], Function]

Returns a zero-argument function that when called returns an instance of cls initialized with the provided function f and name.

The class cls is supposed to be a Function subclass. Calling this method on a suitable Function subclass creates a zero-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances of the function.

Examples:

>>> Phi.make('Proj3', lambda i, e, j: j)
<function Function.make.<locals>.<lambda> at 0x...>

Parameters:

• name (str | None) –

  The name of the function.

• f (Callable) –

  The function that will be used to instantiate cls.

Source code in libmg/compiler/functions.py

@classmethod
def make(cls, name: str | None, f: Callable) -> Callable[[], Function]:
    """Returns a zero-argument function that when called returns an instance of ``cls`` initialized with the provided function ``f`` and ``name``.

    The class ``cls`` is supposed to be a ``Function`` subclass. Calling this method on a suitable ``Function`` subclass
    creates a zero-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary
    automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances
    of the function.

    Examples:
        >>> Phi.make('Proj3', lambda i, e, j: j)
        <function Function.make.<locals>.<lambda> at 0x...>

    Args:
        name: The name of the function.
        f: The function that will be used to instantiate ``cls``.
    """
    if isinstance(f, tf.keras.layers.Layer):  # if f is a layer, regenerate from config to get new weights
        return lambda: cls(type(f).from_config(f.get_config()), name)  # type: ignore
    else:
        return lambda: cls(f, name)

make_parametrized classmethod

make_parametrized(name: str | None, f: Callable[[str], Callable] | Callable[..., Any]) -> Callable[[str], Function]

Returns a one-argument function that when called with the argument a returns an instance of cls initialized with the result of the application of a to the function f and name.

The class cls is supposed to be a Function subclass. Calling this method on a suitable Function subclass creates a one-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances of the function. The function f may have the form lambda x: lambda ... : ... or lambda x, ... : .... The one argument of the lambda returned by this function corresponds to the x argument of f.

Examples:

>>> Phi.make_parametrized('Add', lambda y: lambda i, e, j: j + int(y) * e)
<function Function.make_parametrized.<locals>.<lambda> at 0x...>
>>> Phi.make_parametrized('Add', lambda y, i, e, j: j + int(y) * e)
<function Function.make_parametrized.<locals>.<lambda> at 0x...>

Parameters:

• name (str | None) –

  The name of the function returned by f.

• f (Callable[[str], Callable] | Callable[..., Any]) –

  The function that when applied to some argument a returns the function that will be used to instantiate cls.

Source code in libmg/compiler/functions.py

@classmethod
def make_parametrized(cls, name: str | None, f: Callable[[str], Callable] | Callable[..., Any]) -> Callable[[str], Function]:
    """Returns a one-argument function that when called with the argument ``a``
     returns an instance of ``cls`` initialized with the result of the application of ``a`` to the function ``f`` and ``name``.

    The class ``cls`` is supposed to be a ``Function`` subclass. Calling this method on a suitable ``Function`` subclass
    creates a one-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary
    automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances
    of the function. The function ``f`` may have the form ``lambda x: lambda ... : ...`` or ``lambda x, ... : ...``. The one argument of the lambda returned
    by this function corresponds to the ``x`` argument of ``f``.

    Examples:
        >>> Phi.make_parametrized('Add', lambda y: lambda i, e, j: j + int(y) * e)
        <function Function.make_parametrized.<locals>.<lambda> at 0x...>
        >>> Phi.make_parametrized('Add', lambda y, i, e, j: j + int(y) * e)
        <function Function.make_parametrized.<locals>.<lambda> at 0x...>

    Args:
        name: The name of the function returned by ``f``.
        f: The function that when applied to some argument ``a`` returns the function that will be used to instantiate ``cls``.
    """
    # check if f has only a single argument e.g. lambda x: lambda y: foo(x)(y)
    if f.__code__.co_argcount == 1:
        return lambda a: cls(f(a), name + '_' + a if name is not None else None)
    else:  # e.g. lambda x, y: foo(x)(y)
        return lambda a: cls(partial(f, a), name + '_' + a if name is not None else None)

PsiGlobal

PsiGlobal(single_op: Callable | None = None, multiple_op: Callable | None = None, name: str | None = None)

Bases: PsiNonLocal

A psi function of the mG language that only applies a global transformation of node labels.

A global transformation (i.e. pooling) of node labels \(\psi: T^* \rightarrow U\). For single graph datasets, which use the SingleGraphLoader, only the single_op parameter is necessary. For multiple graph datasets, using the MultipleGraphLoader, only the multiple_op parameter is necessary. The multiple_op argument is a function which takes an additional parameter to distinguish which values in the first argument refer to which graph. For more information, refer to the disjoint data mode in the Spektral library documentation <https://graphneural.network/data-modes/#disjoint-mode/>_.

Examples:

>>> PsiGlobal(single_op=lambda x: tf.reduce_sum(x, axis=0, keepdims=True), multiple_op=lambda x, i: tf.math.segment_sum(x, i), name='SumPooling')
<PsiGlobal ...>

Parameters:

• single_op (Callable | None, default: None ) –

  The function to be used in conjunction with a SingleGraphLoader. The function is expected to take in input a tensor X with shape (n_nodes, n_node_features), containing the labels of every node in the graph, and return a tensor of shape (n_pooled_features, ), containing the pooled output. This output is then broadcast to label every node.

• multiple_op (Callable | None, default: None ) –

  The function to be used in conjunction with a MultipleGraphLoader. The function is expected to take in input a tensor with shape (n_nodes, n_node_features), containing the labels of every node in the graph, and a tensor of graph indices of shape (n_nodes, 1), that mark to which graph every node belongs. The function is expected to return a tensor of shape (n_graphs, n_pooled_features) containing the pooled outputs for each distinct graph index. This output is then broadcast to label every node of each graph with the pooled features for that graph.

• name (str | None, default: None ) –

  The name of the function.

Source code in libmg/compiler/functions.py

def __init__(self, single_op: Callable | None = None,
             multiple_op: Callable | None = None, name: str | None = None):
    """Initializes the instance with a function for a single graph and/or a function a multiple graph operation, and a name.

    Args:
        single_op: The function to be used in conjunction with a ``SingleGraphLoader``.
            The function is expected to take in input a tensor X with shape ``(n_nodes, n_node_features)``, containing the labels of
            every node in the graph, and return a tensor of shape ``(n_pooled_features, )``, containing the pooled output. This output is then broadcast to
            label every node.
        multiple_op: The function to be used in conjunction with a ``MultipleGraphLoader``.
            The function is expected to take in input a tensor with shape ``(n_nodes, n_node_features)``, containing the labels of
            every node in the graph, and a tensor of graph indices of shape ``(n_nodes, 1)``, that mark to which graph every node belongs. The function is
            expected to return a tensor of shape ``(n_graphs, n_pooled_features)`` containing the pooled outputs for each distinct graph index. This output
            is then broadcast to label every node of each graph with the pooled features for that graph.
        name: The name of the function.
    """
    super().__init__(single_op=single_op, multiple_op=multiple_op, name=name)

fname property

fname

The name of this function. This can be either the name provided during initialization, if it was provided, or the dynamic class name.

make classmethod

make(name: str | None, single_op: Callable | None = None, multiple_op: Callable | None = None) -> Callable[[], PsiNonLocal]

Returns a zero-argument function that when called returns an instance of cls initialized with the provided functions single_op and/or multiple_op and name.

Calling this method on a PsiNonLocal class creates a zero-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances of the function.

Examples:

>>> PsiNonLocal.make('Successor', lambda x: x + 1)
<function PsiNonLocal.make.<locals>.<lambda> at 0x...>

Parameters:

• name (str | None) –

  The name of the function returned by single_op and multiple_op.

• single_op (Callable | None, default: None ) –

  The function that will be used to instantiate cls as the single_op argument.

• multiple_op (Callable | None, default: None ) –

  The function that will be used to instantiate cls as the multiple_op argument.

Raises:

• ValueError –

  Neither single_op nor multiple_op have been provided.

Source code in libmg/compiler/functions.py

@classmethod
def make(cls, name: str | None, single_op: Callable | None = None, multiple_op: Callable | None = None) -> Callable[[], PsiNonLocal]:
    """Returns a zero-argument function that when called returns an instance of ``cls`` initialized with the provided functions ``single_op`` and/or
    ``multiple_op`` and ``name``.

    Calling this method on a ``PsiNonLocal`` class creates a zero-argument lambda that returns an instance of such subclass. This way, whenever that
    function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights,
    which are not supposed to be shared with other instances of the function.

    Examples:
        >>> PsiNonLocal.make('Successor', lambda x: x + 1)
        <function PsiNonLocal.make.<locals>.<lambda> at 0x...>

    Args:
        name: The name of the function returned by ``single_op`` and ``multiple_op``.
        single_op: The function that will be used to instantiate ``cls`` as the ``single_op`` argument.
        multiple_op: The function that will be used to instantiate ``cls`` as the ``multiple_op`` argument.

    Raises:
        ValueError: Neither ``single_op`` nor ``multiple_op`` have been provided.
    """
    if single_op is None and multiple_op is None:
        raise ValueError("At least one function must be provided.")
    args = {}
    if single_op is not None:
        args['single_op' if cls is not PsiLocal else 'f'] = single_op
    if multiple_op is not None:
        args['multiple_op'] = multiple_op
    return lambda: cls(**{k: type(v).from_config(v.get_config()) if isinstance(v, tf.keras.layers.Layer) else v for k, v in args.items()}, name=name)

make_parametrized classmethod

make_parametrized(
    name: str | None,
    single_op: Callable[[str], Callable] | Callable[..., Any] | None = None,
    multiple_op: Callable[[str], Callable] | Callable[..., Any] | None = None,
) -> Callable[[str], PsiNonLocal]

Returns a one-argument function that when called with the argument a returns an instance of cls initialized with the result of the application of a to the function single_op and/or multiple_op, and name.

Calling this method on a PsiNonLocal class creates a one-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances of the function. The functions single_op and multiple_op may have the form lambda x: lambda ... : ... or lambda x, ... : .... The one argument of the lambda returned by this function corresponds to the x argument of single_op and multiple_op.

Examples:

>>> PsiNonLocal.make_parametrized(name='Add', single_op=lambda y: lambda x: x + y, multiple_op=lambda y: lambda x, i: x + y)
<function PsiNonLocal.make_parametrized.<locals>.<lambda> at 0x...>
>>> PsiNonLocal.make_parametrized( name='Add', single_op=lambda y, x: x + y, multiple_op=lambda y, x, i: x + y)
<function PsiNonLocal.make_parametrized.<locals>.<lambda> at 0x...>

Parameters:

• name (str | None) –

  The name of the function returned by single_op and multiple_op.

• single_op (Callable[[str], Callable] | Callable[..., Any] | None, default: None ) –

  The function that when applied to some argument a returns the function that will be used to instantiate clsas the single_op argument.

• multiple_op (Callable[[str], Callable] | Callable[..., Any] | None, default: None ) –

  The function that when applied to some argument a returns the function that will be used to instantiate cls as the multiple_op argument.

Raises:

• ValueError –

  Neither single_op nor multiple_op have been provided.

Source code in libmg/compiler/functions.py

@classmethod
def make_parametrized(cls, name: str | None, single_op: Callable[[str], Callable] | Callable[..., Any] | None = None,
                      multiple_op: Callable[[str], Callable] | Callable[..., Any] | None = None) -> Callable[[str], PsiNonLocal]:
    """Returns a one-argument function that when called with the argument ``a`` returns an instance of ``cls`` initialized with the result of the
     application of ``a`` to the function ``single_op`` and/or ``multiple_op``, and ``name``.

    Calling this method on a ``PsiNonLocal`` class creates a one-argument lambda that returns an instance of such subclass. This way, whenever that
    function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights,
    which are not supposed to be shared with other instances of the function. The functions ``single_op`` and ``multiple_op`` may have the form ``lambda
    x: lambda ... : ...`` or ``lambda x, ... : ...``. The one argument of the lambda returned by this function corresponds to the ``x`` argument of
    ``single_op`` and ``multiple_op``.

    Examples:
        >>> PsiNonLocal.make_parametrized(name='Add', single_op=lambda y: lambda x: x + y, multiple_op=lambda y: lambda x, i: x + y)
        <function PsiNonLocal.make_parametrized.<locals>.<lambda> at 0x...>
        >>> PsiNonLocal.make_parametrized( name='Add', single_op=lambda y, x: x + y, multiple_op=lambda y, x, i: x + y)
        <function PsiNonLocal.make_parametrized.<locals>.<lambda> at 0x...>

    Args:
        name: The name of the function returned by ``single_op`` and ``multiple_op``.
        single_op: The function that when applied to some argument ``a`` returns the function that will be used to instantiate ``cls``as the
            ``single_op`` argument.
        multiple_op: The function that when applied to some argument ``a`` returns the function that will be used to instantiate ``cls`` as the
            ``multiple_op`` argument.

    Raises:
        ValueError: Neither ``single_op`` nor ``multiple_op`` have been provided.
    """
    if single_op is None and multiple_op is None:
        raise ValueError("At least one function must be provided.")
    args = {}
    if single_op is not None:
        args['single_op' if cls is not PsiLocal else 'f'] = single_op
    if multiple_op is not None:
        args['multiple_op'] = multiple_op
    return lambda a: cls(**{k: v(a) if v.__code__.co_argcount == 1 else partial(v, a) for k, v in args.items()},
                         name=name + '_' + a if name is not None else None)

PsiNonLocal

PsiNonLocal(single_op: Callable | None = None, multiple_op: Callable | None = None, name: str | None = None)

Bases: Psi

A psi function of the mG language.

A non-local function applied on node labels \(\psi: T^* \times T \rightarrow U\). For single graph datasets, which use the SingleGraphLoader, only the single_op parameter is necessary. For multiple graph datasets, using the MultipleGraphLoader, only the multiple_op parameter is necessary. The multiple_op argument is a function which takes an additional parameter to distinguish which values in the first argument refer to which graph. For more information, refer to the disjoint data mode in the Spektral library documentation <https://graphneural.network/data-modes/#disjoint-mode/>_.

Examples:

>>> PsiNonLocal(single_op=lambda x: x + 1, multiple_op=lambda x, i: x + 1, name='Add1')
<PsiNonLocal ...>

Attributes:

• single_op –

  A function to be used in conjunction with a SingleGraphLoader

• multiple_op –

  A function to be used in conjunction with a MultipleGraphLoader

Parameters:

• single_op (Callable | None, default: None ) –

  The function to be used in conjunction with a SingleGraphLoader. The function must be compatible with TensorFlow's broadcasting rules. The function is expected to take in input a tensor X with shape (n_nodes, n_node_features), containing the labels of every node in the graph, and return a tensor of shape (n_nodes, n_new_node_features), containing the transformed labels. The function can use broadcasting to emulate the tuple (T*, T) in the definition of psi.

• multiple_op (Callable | None, default: None ) –

  The function to be used in conjunction with a MultipleGraphLoader. The function must be compatible with TensorFlow's broadcasting rules. The function is expected to take in input a tensor X with shape (n_nodes, n_node_features), containing the labels of every node in the graph, and a tensor of graph indices of shape (n_nodes, 1), that mark to which graph every node belongs. The function is expected to return a tensor of shape (n_nodes, n_new_node_features) containing the transformed labels. The function can use broadcasting to emulate the tuple (T*, T) in the definition of psi.

• name (str | None, default: None ) –

  The name of the function.

Source code in libmg/compiler/functions.py

def __init__(self, single_op: Callable | None = None,
             multiple_op: Callable | None = None,
             name: str | None = None):
    """Initializes the instance with a function for a single graph and/or a function a multiple graph operation, and a name.

    Args:
        single_op: The function to be used in conjunction with a ``SingleGraphLoader``. The function must be compatible with TensorFlow's broadcasting
            rules. The function is expected to take in input a tensor X with shape ``(n_nodes, n_node_features)``, containing the labels of
            every node in the graph, and return a tensor of shape ``(n_nodes, n_new_node_features)``, containing the transformed labels.
            The function can use broadcasting to emulate the tuple (T*, T) in the definition of psi.
        multiple_op: The function to be used in conjunction with a ``MultipleGraphLoader``. The function must be compatible with TensorFlow's broadcasting
            rules. The function is expected to take in input a tensor X with shape ``(n_nodes, n_node_features)``, containing the labels of
            every node in the graph, and a tensor of graph indices of shape ``(n_nodes, 1)``, that mark to which graph every node belongs. The function is
            expected to return a tensor of shape ``(n_nodes, n_new_node_features)`` containing the transformed labels.
            The function can use broadcasting to emulate the tuple (T*, T) in the definition of psi.
        name: The name of the function.
    """
    if single_op is None:
        self.single_op = self.single_graph_op
    else:
        self.single_op = single_op  # type: ignore
    if multiple_op is None:
        self.multiple_op = self.multiple_graph_op
    else:
        self.multiple_op = multiple_op  # type: ignore
    super().__init__(self.single_op, name)

fname property

fname

The name of this function. This can be either the name provided during initialization, if it was provided, or the dynamic class name.

make classmethod

make(name: str | None, single_op: Callable | None = None, multiple_op: Callable | None = None) -> Callable[[], PsiNonLocal]

Returns a zero-argument function that when called returns an instance of cls initialized with the provided functions single_op and/or multiple_op and name.

Calling this method on a PsiNonLocal class creates a zero-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances of the function.

Examples:

>>> PsiNonLocal.make('Successor', lambda x: x + 1)
<function PsiNonLocal.make.<locals>.<lambda> at 0x...>

Parameters:

• name (str | None) –

  The name of the function returned by single_op and multiple_op.

• single_op (Callable | None, default: None ) –

  The function that will be used to instantiate cls as the single_op argument.

• multiple_op (Callable | None, default: None ) –

  The function that will be used to instantiate cls as the multiple_op argument.

Raises:

• ValueError –

  Neither single_op nor multiple_op have been provided.

Source code in libmg/compiler/functions.py

@classmethod
def make(cls, name: str | None, single_op: Callable | None = None, multiple_op: Callable | None = None) -> Callable[[], PsiNonLocal]:
    """Returns a zero-argument function that when called returns an instance of ``cls`` initialized with the provided functions ``single_op`` and/or
    ``multiple_op`` and ``name``.

    Calling this method on a ``PsiNonLocal`` class creates a zero-argument lambda that returns an instance of such subclass. This way, whenever that
    function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights,
    which are not supposed to be shared with other instances of the function.

    Examples:
        >>> PsiNonLocal.make('Successor', lambda x: x + 1)
        <function PsiNonLocal.make.<locals>.<lambda> at 0x...>

    Args:
        name: The name of the function returned by ``single_op`` and ``multiple_op``.
        single_op: The function that will be used to instantiate ``cls`` as the ``single_op`` argument.
        multiple_op: The function that will be used to instantiate ``cls`` as the ``multiple_op`` argument.

    Raises:
        ValueError: Neither ``single_op`` nor ``multiple_op`` have been provided.
    """
    if single_op is None and multiple_op is None:
        raise ValueError("At least one function must be provided.")
    args = {}
    if single_op is not None:
        args['single_op' if cls is not PsiLocal else 'f'] = single_op
    if multiple_op is not None:
        args['multiple_op'] = multiple_op
    return lambda: cls(**{k: type(v).from_config(v.get_config()) if isinstance(v, tf.keras.layers.Layer) else v for k, v in args.items()}, name=name)

make_parametrized classmethod

make_parametrized(
    name: str | None,
    single_op: Callable[[str], Callable] | Callable[..., Any] | None = None,
    multiple_op: Callable[[str], Callable] | Callable[..., Any] | None = None,
) -> Callable[[str], PsiNonLocal]

Returns a one-argument function that when called with the argument a returns an instance of cls initialized with the result of the application of a to the function single_op and/or multiple_op, and name.

Calling this method on a PsiNonLocal class creates a one-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances of the function. The functions single_op and multiple_op may have the form lambda x: lambda ... : ... or lambda x, ... : .... The one argument of the lambda returned by this function corresponds to the x argument of single_op and multiple_op.

Examples:

>>> PsiNonLocal.make_parametrized(name='Add', single_op=lambda y: lambda x: x + y, multiple_op=lambda y: lambda x, i: x + y)
<function PsiNonLocal.make_parametrized.<locals>.<lambda> at 0x...>
>>> PsiNonLocal.make_parametrized( name='Add', single_op=lambda y, x: x + y, multiple_op=lambda y, x, i: x + y)
<function PsiNonLocal.make_parametrized.<locals>.<lambda> at 0x...>

Parameters:

• name (str | None) –

  The name of the function returned by single_op and multiple_op.

• single_op (Callable[[str], Callable] | Callable[..., Any] | None, default: None ) –

  The function that when applied to some argument a returns the function that will be used to instantiate clsas the single_op argument.

• multiple_op (Callable[[str], Callable] | Callable[..., Any] | None, default: None ) –

  The function that when applied to some argument a returns the function that will be used to instantiate cls as the multiple_op argument.

Raises:

• ValueError –

  Neither single_op nor multiple_op have been provided.

Source code in libmg/compiler/functions.py

@classmethod
def make_parametrized(cls, name: str | None, single_op: Callable[[str], Callable] | Callable[..., Any] | None = None,
                      multiple_op: Callable[[str], Callable] | Callable[..., Any] | None = None) -> Callable[[str], PsiNonLocal]:
    """Returns a one-argument function that when called with the argument ``a`` returns an instance of ``cls`` initialized with the result of the
     application of ``a`` to the function ``single_op`` and/or ``multiple_op``, and ``name``.

    Calling this method on a ``PsiNonLocal`` class creates a one-argument lambda that returns an instance of such subclass. This way, whenever that
    function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights,
    which are not supposed to be shared with other instances of the function. The functions ``single_op`` and ``multiple_op`` may have the form ``lambda
    x: lambda ... : ...`` or ``lambda x, ... : ...``. The one argument of the lambda returned by this function corresponds to the ``x`` argument of
    ``single_op`` and ``multiple_op``.

    Examples:
        >>> PsiNonLocal.make_parametrized(name='Add', single_op=lambda y: lambda x: x + y, multiple_op=lambda y: lambda x, i: x + y)
        <function PsiNonLocal.make_parametrized.<locals>.<lambda> at 0x...>
        >>> PsiNonLocal.make_parametrized( name='Add', single_op=lambda y, x: x + y, multiple_op=lambda y, x, i: x + y)
        <function PsiNonLocal.make_parametrized.<locals>.<lambda> at 0x...>

    Args:
        name: The name of the function returned by ``single_op`` and ``multiple_op``.
        single_op: The function that when applied to some argument ``a`` returns the function that will be used to instantiate ``cls``as the
            ``single_op`` argument.
        multiple_op: The function that when applied to some argument ``a`` returns the function that will be used to instantiate ``cls`` as the
            ``multiple_op`` argument.

    Raises:
        ValueError: Neither ``single_op`` nor ``multiple_op`` have been provided.
    """
    if single_op is None and multiple_op is None:
        raise ValueError("At least one function must be provided.")
    args = {}
    if single_op is not None:
        args['single_op' if cls is not PsiLocal else 'f'] = single_op
    if multiple_op is not None:
        args['multiple_op'] = multiple_op
    return lambda a: cls(**{k: v(a) if v.__code__.co_argcount == 1 else partial(v, a) for k, v in args.items()},
                         name=name + '_' + a if name is not None else None)

Phi

Phi(f: Callable[[tuple[tf.Tensor, ...], tf.Tensor, tuple[tf.Tensor, ...]], tf.Tensor] | None = None, name: str | None = None)

Bases: Function

A phi function of the mG language.

A function \(\varphi: T \times U \times T \rightarrow V\) to compute the message sent by a node i to a node j through edge e.

Examples:

>>> Phi(lambda i, e, j: i * e, name='EdgeProd')
<Phi ...>

Parameters:

• f (Callable[[tuple[Tensor, ...], Tensor, tuple[Tensor, ...]], Tensor] | None, default: None ) –

  The function that this object will run when called. The function must be compatible with TensorFlow's broadcasting rules. The function is expected to take in input a tensor X1 with shape (n_edges, n_node_features), containing the labels of all nodes sending a message, a tensor E with shape (n_edges, n_edge_features), containing the labels of all edges in the graph, and a tensor X2, containing the labels of all nodes receiving a message. The function is expected to return a tensor with shape (n_edges, n_message_features), containing the messages to be sent to the destination nodes.

• name (str | None, default: None ) –

  The name of the function.

Source code in libmg/compiler/functions.py

def __init__(self, f: Callable[[tuple[tf.Tensor, ...], tf.Tensor, tuple[tf.Tensor, ...]], tf.Tensor] | None = None,
             name: str | None = None):
    """Initializes the instance with a function and a name.

    Args:
        f: The function that this object will run when called. The function must be compatible with TensorFlow's broadcasting
            rules. The function is expected to take in input a tensor X1 with shape ``(n_edges, n_node_features)``, containing the labels of
            all nodes sending a message, a tensor E with shape ``(n_edges, n_edge_features)``, containing the labels of all edges in the graph, and a tensor
            X2, containing the labels of all nodes receiving a message. The function is expected to return a tensor with shape
            ``(n_edges, n_message_features)``, containing the messages to be sent to the destination nodes.
        name: The name of the function.
    """
    if f is None:
        f = self.func
    super().__init__(f, name)

fname property

fname

The name of this function. This can be either the name provided during initialization, if it was provided, or the dynamic class name.

make classmethod

make(name: str | None, f: Callable) -> Callable[[], Function]

Returns a zero-argument function that when called returns an instance of cls initialized with the provided function f and name.

The class cls is supposed to be a Function subclass. Calling this method on a suitable Function subclass creates a zero-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances of the function.

Examples:

>>> Phi.make('Proj3', lambda i, e, j: j)
<function Function.make.<locals>.<lambda> at 0x...>

Parameters:

• name (str | None) –

  The name of the function.

• f (Callable) –

  The function that will be used to instantiate cls.

Source code in libmg/compiler/functions.py

@classmethod
def make(cls, name: str | None, f: Callable) -> Callable[[], Function]:
    """Returns a zero-argument function that when called returns an instance of ``cls`` initialized with the provided function ``f`` and ``name``.

    The class ``cls`` is supposed to be a ``Function`` subclass. Calling this method on a suitable ``Function`` subclass
    creates a zero-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary
    automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances
    of the function.

    Examples:
        >>> Phi.make('Proj3', lambda i, e, j: j)
        <function Function.make.<locals>.<lambda> at 0x...>

    Args:
        name: The name of the function.
        f: The function that will be used to instantiate ``cls``.
    """
    if isinstance(f, tf.keras.layers.Layer):  # if f is a layer, regenerate from config to get new weights
        return lambda: cls(type(f).from_config(f.get_config()), name)  # type: ignore
    else:
        return lambda: cls(f, name)

make_parametrized classmethod

make_parametrized(name: str | None, f: Callable[[str], Callable] | Callable[..., Any]) -> Callable[[str], Function]

Returns a one-argument function that when called with the argument a returns an instance of cls initialized with the result of the application of a to the function f and name.

The class cls is supposed to be a Function subclass. Calling this method on a suitable Function subclass creates a one-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances of the function. The function f may have the form lambda x: lambda ... : ... or lambda x, ... : .... The one argument of the lambda returned by this function corresponds to the x argument of f.

Examples:

>>> Phi.make_parametrized('Add', lambda y: lambda i, e, j: j + int(y) * e)
<function Function.make_parametrized.<locals>.<lambda> at 0x...>
>>> Phi.make_parametrized('Add', lambda y, i, e, j: j + int(y) * e)
<function Function.make_parametrized.<locals>.<lambda> at 0x...>

Parameters:

• name (str | None) –

  The name of the function returned by f.

• f (Callable[[str], Callable] | Callable[..., Any]) –

  The function that when applied to some argument a returns the function that will be used to instantiate cls.

Source code in libmg/compiler/functions.py

@classmethod
def make_parametrized(cls, name: str | None, f: Callable[[str], Callable] | Callable[..., Any]) -> Callable[[str], Function]:
    """Returns a one-argument function that when called with the argument ``a``
     returns an instance of ``cls`` initialized with the result of the application of ``a`` to the function ``f`` and ``name``.

    The class ``cls`` is supposed to be a ``Function`` subclass. Calling this method on a suitable ``Function`` subclass
    creates a one-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary
    automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances
    of the function. The function ``f`` may have the form ``lambda x: lambda ... : ...`` or ``lambda x, ... : ...``. The one argument of the lambda returned
    by this function corresponds to the ``x`` argument of ``f``.

    Examples:
        >>> Phi.make_parametrized('Add', lambda y: lambda i, e, j: j + int(y) * e)
        <function Function.make_parametrized.<locals>.<lambda> at 0x...>
        >>> Phi.make_parametrized('Add', lambda y, i, e, j: j + int(y) * e)
        <function Function.make_parametrized.<locals>.<lambda> at 0x...>

    Args:
        name: The name of the function returned by ``f``.
        f: The function that when applied to some argument ``a`` returns the function that will be used to instantiate ``cls``.
    """
    # check if f has only a single argument e.g. lambda x: lambda y: foo(x)(y)
    if f.__code__.co_argcount == 1:
        return lambda a: cls(f(a), name + '_' + a if name is not None else None)
    else:  # e.g. lambda x, y: foo(x)(y)
        return lambda a: cls(partial(f, a), name + '_' + a if name is not None else None)

Sigma

Sigma(f: Callable[[tf.Tensor, tf.Tensor, int, tuple[tf.Tensor, ...]], tf.Tensor] | None = None, name: str | None = None)

Bases: Function

A sigma function of the mG language.

A function \(\sigma: T^* \times U \rightarrow V\) to aggregate the messages sent to a node, including the current label of the node.

Examples:

>>> Sigma(lambda m, i, n, x: tf.math.segment_max(m, i), name='Max')
<Sigma ...>

Parameters:

• f (Callable[[Tensor, Tensor, int, tuple[Tensor, ...]], Tensor] | None, default: None ) –

  The function that this object will run when called. The function must be compatible with TensorFlow's broadcasting rules. The function is expected to take in input a tensor M with shape (n_edges, n_message_features), containing the generated messages, a tensor IDX with shape (n_edges,), containing the ids of the node each message is being sent to, the total number of the nodes involved, and a tensor X, containing the current node labels. The function is expected to return a tensor with shape (n_nodes, n_new_node_features), containing the new node labels.

• name (str | None, default: None ) –

  The name of the function.

Source code in libmg/compiler/functions.py

def __init__(self, f: Callable[[tf.Tensor, tf.Tensor, int, tuple[tf.Tensor, ...]], tf.Tensor] | None = None,
             name: str | None = None):
    """Initializes the instance with a function and a name.

    Args:
        f: The function that this object will run when called. The function must be compatible with TensorFlow's broadcasting
            rules. The function is expected to take in input a tensor M with shape ``(n_edges, n_message_features)``, containing the generated messages,
            a tensor IDX with shape ``(n_edges,)``, containing the ids of the node each message is being sent to, the total number of the nodes involved,
            and a tensor X, containing the current node labels. The function is expected to return a tensor with shape
            ``(n_nodes, n_new_node_features)``, containing the new node labels.
        name: The name of the function.
    """
    if f is None:
        f = self.func
    super().__init__(f, name)

fname property

fname

The name of this function. This can be either the name provided during initialization, if it was provided, or the dynamic class name.

make classmethod

make(name: str | None, f: Callable) -> Callable[[], Function]

Returns a zero-argument function that when called returns an instance of cls initialized with the provided function f and name.

The class cls is supposed to be a Function subclass. Calling this method on a suitable Function subclass creates a zero-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances of the function.

Examples:

>>> Phi.make('Proj3', lambda i, e, j: j)
<function Function.make.<locals>.<lambda> at 0x...>

Parameters:

• name (str | None) –

  The name of the function.

• f (Callable) –

  The function that will be used to instantiate cls.

Source code in libmg/compiler/functions.py

@classmethod
def make(cls, name: str | None, f: Callable) -> Callable[[], Function]:
    """Returns a zero-argument function that when called returns an instance of ``cls`` initialized with the provided function ``f`` and ``name``.

    The class ``cls`` is supposed to be a ``Function`` subclass. Calling this method on a suitable ``Function`` subclass
    creates a zero-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary
    automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances
    of the function.

    Examples:
        >>> Phi.make('Proj3', lambda i, e, j: j)
        <function Function.make.<locals>.<lambda> at 0x...>

    Args:
        name: The name of the function.
        f: The function that will be used to instantiate ``cls``.
    """
    if isinstance(f, tf.keras.layers.Layer):  # if f is a layer, regenerate from config to get new weights
        return lambda: cls(type(f).from_config(f.get_config()), name)  # type: ignore
    else:
        return lambda: cls(f, name)

make_parametrized classmethod

make_parametrized(name: str | None, f: Callable[[str], Callable] | Callable[..., Any]) -> Callable[[str], Function]

Returns a one-argument function that when called with the argument a returns an instance of cls initialized with the result of the application of a to the function f and name.

The class cls is supposed to be a Function subclass. Calling this method on a suitable Function subclass creates a one-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances of the function. The function f may have the form lambda x: lambda ... : ... or lambda x, ... : .... The one argument of the lambda returned by this function corresponds to the x argument of f.

Examples:

>>> Phi.make_parametrized('Add', lambda y: lambda i, e, j: j + int(y) * e)
<function Function.make_parametrized.<locals>.<lambda> at 0x...>
>>> Phi.make_parametrized('Add', lambda y, i, e, j: j + int(y) * e)
<function Function.make_parametrized.<locals>.<lambda> at 0x...>

Parameters:

• name (str | None) –

  The name of the function returned by f.

• f (Callable[[str], Callable] | Callable[..., Any]) –

  The function that when applied to some argument a returns the function that will be used to instantiate cls.

Source code in libmg/compiler/functions.py

@classmethod
def make_parametrized(cls, name: str | None, f: Callable[[str], Callable] | Callable[..., Any]) -> Callable[[str], Function]:
    """Returns a one-argument function that when called with the argument ``a``
     returns an instance of ``cls`` initialized with the result of the application of ``a`` to the function ``f`` and ``name``.

    The class ``cls`` is supposed to be a ``Function`` subclass. Calling this method on a suitable ``Function`` subclass
    creates a one-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary
    automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances
    of the function. The function ``f`` may have the form ``lambda x: lambda ... : ...`` or ``lambda x, ... : ...``. The one argument of the lambda returned
    by this function corresponds to the ``x`` argument of ``f``.

    Examples:
        >>> Phi.make_parametrized('Add', lambda y: lambda i, e, j: j + int(y) * e)
        <function Function.make_parametrized.<locals>.<lambda> at 0x...>
        >>> Phi.make_parametrized('Add', lambda y, i, e, j: j + int(y) * e)
        <function Function.make_parametrized.<locals>.<lambda> at 0x...>

    Args:
        name: The name of the function returned by ``f``.
        f: The function that when applied to some argument ``a`` returns the function that will be used to instantiate ``cls``.
    """
    # check if f has only a single argument e.g. lambda x: lambda y: foo(x)(y)
    if f.__code__.co_argcount == 1:
        return lambda a: cls(f(a), name + '_' + a if name is not None else None)
    else:  # e.g. lambda x, y: foo(x)(y)
        return lambda a: cls(partial(f, a), name + '_' + a if name is not None else None)

Constant

Constant(v: tf.Tensor, name: str | None = None)

Bases: PsiLocal

A constant psi function of the mG language.

A constant function \(\psi: T \rightarrow U\) that maps every node label to a constant value.

Examples:

>>> Constant(tf.constant(False), name='False')
<Constant ...>

Parameters:

• v (Tensor) –

  A scalar or tensor value that identifies the constant function.

• name (str | None, default: None ) –

  The name of the function.

Source code in libmg/compiler/functions.py

def __init__(self, v: tf.Tensor, name: str | None = None):
    """Initializes the instance with a function and a name.

    Args:
        v: A scalar or tensor value that identifies the constant function.
        name: The name of the function.
    """
    if tf.reduce_all(tf.equal(tf.cast(v, dtype=tf.float32), 1.0)):
        def f(x):
            return tf.ones((tf.shape(x)[0], tf.size(v)), dtype=v.dtype)
    elif tf.reduce_all(tf.equal(tf.cast(v, dtype=tf.float32), 0.0)):
        def f(x):
            return tf.zeros((tf.shape(x)[0], tf.size(v)), dtype=v.dtype)
    elif tf.size(v) == 1:
        def f(x):
            return tf.fill((tf.shape(x)[0], tf.size(v)), value=v)
    else:
        def f(x):
            return tf.tile([v], [tf.shape(x)[0], 1])
    super().__init__(f, name)

fname property

fname

The name of this function. This can be either the name provided during initialization, if it was provided, or the dynamic class name.

make classmethod

make(name: str | None, f: Callable) -> Callable[[], Function]

Returns a zero-argument function that when called returns an instance of cls initialized with the provided function f and name.

The class cls is supposed to be a Function subclass. Calling this method on a suitable Function subclass creates a zero-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances of the function.

Examples:

>>> Phi.make('Proj3', lambda i, e, j: j)
<function Function.make.<locals>.<lambda> at 0x...>

Parameters:

• name (str | None) –

  The name of the function.

• f (Callable) –

  The function that will be used to instantiate cls.

Source code in libmg/compiler/functions.py

@classmethod
def make(cls, name: str | None, f: Callable) -> Callable[[], Function]:
    """Returns a zero-argument function that when called returns an instance of ``cls`` initialized with the provided function ``f`` and ``name``.

    The class ``cls`` is supposed to be a ``Function`` subclass. Calling this method on a suitable ``Function`` subclass
    creates a zero-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary
    automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances
    of the function.

    Examples:
        >>> Phi.make('Proj3', lambda i, e, j: j)
        <function Function.make.<locals>.<lambda> at 0x...>

    Args:
        name: The name of the function.
        f: The function that will be used to instantiate ``cls``.
    """
    if isinstance(f, tf.keras.layers.Layer):  # if f is a layer, regenerate from config to get new weights
        return lambda: cls(type(f).from_config(f.get_config()), name)  # type: ignore
    else:
        return lambda: cls(f, name)

make_parametrized classmethod

make_parametrized(name: str | None, f: Callable[[str], Callable] | Callable[..., Any]) -> Callable[[str], Function]

Returns a one-argument function that when called with the argument a returns an instance of cls initialized with the result of the application of a to the function f and name.

The class cls is supposed to be a Function subclass. Calling this method on a suitable Function subclass creates a one-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances of the function. The function f may have the form lambda x: lambda ... : ... or lambda x, ... : .... The one argument of the lambda returned by this function corresponds to the x argument of f.

Examples:

>>> Phi.make_parametrized('Add', lambda y: lambda i, e, j: j + int(y) * e)
<function Function.make_parametrized.<locals>.<lambda> at 0x...>
>>> Phi.make_parametrized('Add', lambda y, i, e, j: j + int(y) * e)
<function Function.make_parametrized.<locals>.<lambda> at 0x...>

Parameters:

• name (str | None) –

  The name of the function returned by f.

• f (Callable[[str], Callable] | Callable[..., Any]) –

  The function that when applied to some argument a returns the function that will be used to instantiate cls.

Source code in libmg/compiler/functions.py

@classmethod
def make_parametrized(cls, name: str | None, f: Callable[[str], Callable] | Callable[..., Any]) -> Callable[[str], Function]:
    """Returns a one-argument function that when called with the argument ``a``
     returns an instance of ``cls`` initialized with the result of the application of ``a`` to the function ``f`` and ``name``.

    The class ``cls`` is supposed to be a ``Function`` subclass. Calling this method on a suitable ``Function`` subclass
    creates a one-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary
    automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances
    of the function. The function ``f`` may have the form ``lambda x: lambda ... : ...`` or ``lambda x, ... : ...``. The one argument of the lambda returned
    by this function corresponds to the ``x`` argument of ``f``.

    Examples:
        >>> Phi.make_parametrized('Add', lambda y: lambda i, e, j: j + int(y) * e)
        <function Function.make_parametrized.<locals>.<lambda> at 0x...>
        >>> Phi.make_parametrized('Add', lambda y, i, e, j: j + int(y) * e)
        <function Function.make_parametrized.<locals>.<lambda> at 0x...>

    Args:
        name: The name of the function returned by ``f``.
        f: The function that when applied to some argument ``a`` returns the function that will be used to instantiate ``cls``.
    """
    # check if f has only a single argument e.g. lambda x: lambda y: foo(x)(y)
    if f.__code__.co_argcount == 1:
        return lambda a: cls(f(a), name + '_' + a if name is not None else None)
    else:  # e.g. lambda x, y: foo(x)(y)
        return lambda a: cls(partial(f, a), name + '_' + a if name is not None else None)

Pi

Pi(i: int, j: int | None = None, name: str | None = None)

Bases: PsiLocal

A projection psi function of the mG language.

A projection function \(\psi: T^n \rightarrow T^m\) that maps every node label to a projection of itself.

Examples:

>>> Pi(0, 2, name='FirstTwo')
<Pi ...>

Parameters:

• i (int) –

  0-based index, start position of the projection function, inclusive.

• j (int | None, default: None ) –

  end position of the projection function, exclusive. Defaults to i + 1.

• name (str | None, default: None ) –

  The name of the function.

Raises:

• ValueError –

  start and end index are equal.

Source code in libmg/compiler/functions.py

def __init__(self, i: int, j: int | None = None, name: str | None = None):
    """Initializes the instance with the projection indexes and a name.

    Args:
        i: 0-based index, start position of the projection function, inclusive.
        j: end position of the projection function, exclusive. Defaults to i + 1.
        name: The name of the function.

    Raises:
        ValueError: start and end index are equal.
    """
    j = i + 1 if j is None else j
    if i == j:
        raise ValueError("Start index and end index cannot be equal.")

    def f(x): return x[:, i:j]

    super().__init__(f, name)

fname property

fname

The name of this function. This can be either the name provided during initialization, if it was provided, or the dynamic class name.

make classmethod

make(name: str | None, f: Callable) -> Callable[[], Function]

Returns a zero-argument function that when called returns an instance of cls initialized with the provided function f and name.

The class cls is supposed to be a Function subclass. Calling this method on a suitable Function subclass creates a zero-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances of the function.

Examples:

>>> Phi.make('Proj3', lambda i, e, j: j)
<function Function.make.<locals>.<lambda> at 0x...>

Parameters:

• name (str | None) –

  The name of the function.

• f (Callable) –

  The function that will be used to instantiate cls.

Source code in libmg/compiler/functions.py

@classmethod
def make(cls, name: str | None, f: Callable) -> Callable[[], Function]:
    """Returns a zero-argument function that when called returns an instance of ``cls`` initialized with the provided function ``f`` and ``name``.

    The class ``cls`` is supposed to be a ``Function`` subclass. Calling this method on a suitable ``Function`` subclass
    creates a zero-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary
    automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances
    of the function.

    Examples:
        >>> Phi.make('Proj3', lambda i, e, j: j)
        <function Function.make.<locals>.<lambda> at 0x...>

    Args:
        name: The name of the function.
        f: The function that will be used to instantiate ``cls``.
    """
    if isinstance(f, tf.keras.layers.Layer):  # if f is a layer, regenerate from config to get new weights
        return lambda: cls(type(f).from_config(f.get_config()), name)  # type: ignore
    else:
        return lambda: cls(f, name)

make_parametrized classmethod

make_parametrized(name: str | None, f: Callable[[str], Callable] | Callable[..., Any]) -> Callable[[str], Function]

Returns a one-argument function that when called with the argument a returns an instance of cls initialized with the result of the application of a to the function f and name.

The class cls is supposed to be a Function subclass. Calling this method on a suitable Function subclass creates a one-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances of the function. The function f may have the form lambda x: lambda ... : ... or lambda x, ... : .... The one argument of the lambda returned by this function corresponds to the x argument of f.

Examples:

>>> Phi.make_parametrized('Add', lambda y: lambda i, e, j: j + int(y) * e)
<function Function.make_parametrized.<locals>.<lambda> at 0x...>
>>> Phi.make_parametrized('Add', lambda y, i, e, j: j + int(y) * e)
<function Function.make_parametrized.<locals>.<lambda> at 0x...>

Parameters:

• name (str | None) –

  The name of the function returned by f.

• f (Callable[[str], Callable] | Callable[..., Any]) –

  The function that when applied to some argument a returns the function that will be used to instantiate cls.

Source code in libmg/compiler/functions.py

@classmethod
def make_parametrized(cls, name: str | None, f: Callable[[str], Callable] | Callable[..., Any]) -> Callable[[str], Function]:
    """Returns a one-argument function that when called with the argument ``a``
     returns an instance of ``cls`` initialized with the result of the application of ``a`` to the function ``f`` and ``name``.

    The class ``cls`` is supposed to be a ``Function`` subclass. Calling this method on a suitable ``Function`` subclass
    creates a one-argument lambda that returns an instance of such subclass. This way, whenever that function is used in a mG program, the dictionary
    automatically regenerates the instance. This is useful when the function has trainable weights, which are not supposed to be shared with other instances
    of the function. The function ``f`` may have the form ``lambda x: lambda ... : ...`` or ``lambda x, ... : ...``. The one argument of the lambda returned
    by this function corresponds to the ``x`` argument of ``f``.

    Examples:
        >>> Phi.make_parametrized('Add', lambda y: lambda i, e, j: j + int(y) * e)
        <function Function.make_parametrized.<locals>.<lambda> at 0x...>
        >>> Phi.make_parametrized('Add', lambda y, i, e, j: j + int(y) * e)
        <function Function.make_parametrized.<locals>.<lambda> at 0x...>

    Args:
        name: The name of the function returned by ``f``.
        f: The function that when applied to some argument ``a`` returns the function that will be used to instantiate ``cls``.
    """
    # check if f has only a single argument e.g. lambda x: lambda y: foo(x)(y)
    if f.__code__.co_argcount == 1:
        return lambda a: cls(f(a), name + '_' + a if name is not None else None)
    else:  # e.g. lambda x, y: foo(x)(y)
        return lambda a: cls(partial(f, a), name + '_' + a if name is not None else None)

NodeConfig

NodeConfig(node_type: tf.DType, node_size: int)

Bases: LabelConfig

Defines the signature of a node label.

Parameters:

• node_type (DType) –

  Type of the node labels.

• node_size (int) –

  Dimension of the node labels.

Source code in libmg/compiler/compiler.py

def __init__(self, node_type: tf.DType, node_size: int):
    """Initializes the instance with the given type and dimension.

    Args:
        node_type: Type of the node labels.
        node_size: Dimension of the node labels.
    """
    super().__init__(node_type, node_size)

EdgeConfig

EdgeConfig(edge_type: tf.DType, edge_size: int)

Bases: LabelConfig

Defines the signature of an edge label.

Parameters:

• edge_type (DType) –

  Type of the edge labels.

• edge_size (int) –

  Dimension of the edge labels.

Source code in libmg/compiler/compiler.py

def __init__(self, edge_type: tf.DType, edge_size: int):
    """Initializes the instance with the given type and dimension.

    Args:
        edge_type: Type of the edge labels.
        edge_size: Dimension of the edge labels.
    """
    super().__init__(edge_type, edge_size)

CompilerConfig

CompilerConfig(node_config: NodeConfig, edge_config: EdgeConfig | None, matrix_type: tf.DType, tolerance: dict[str, float], multiple_loader: bool)

Defines the configuration for the mG compiler.

It is recommended to use the static constructor methods to instantiate this class.

Parameters:

• node_config (NodeConfig) –

  The signature of the initial node labels of the graphs.

• edge_config (EdgeConfig | None) –

  The signature of the initial edge labels of the graphs.

• matrix_type (DType) –

  The signature of the adjacency matrix of the graphs.

• tolerance (dict[str, float]) –

  The tolerance values for the data types (typically floats) that require an approximate fixed point solution.

• multiple_loader (bool) –

  True if the models generated by the mG compiler will receive their inputs by a MultipleGraphLoader.

Source code in libmg/compiler/compiler.py

def __init__(self, node_config: NodeConfig, edge_config: EdgeConfig | None, matrix_type: tf.DType, tolerance: dict[str, float], multiple_loader: bool):
    """Initializes the instance with the initial signature of the node labels and, if present, edge labels, the signature of the adjacency matrix, the
    tolerance values by data type and which loader will be used for the graphs.

    Args:
        node_config: The signature of the initial node labels of the graphs.
        edge_config: The signature of the initial edge labels of the graphs.
        matrix_type: The signature of the adjacency matrix of the graphs.
        tolerance: The tolerance values for the data types (typically floats) that require an approximate fixed point solution.
        multiple_loader: True if the models generated by the mG compiler will receive their inputs by a ``MultipleGraphLoader``.
    """
    self._node_config = node_config
    self._edge_config = edge_config
    self._matrix_type = matrix_type
    self._tolerance = tolerance
    self._multiple_loader = multiple_loader

node_feature_type property

node_feature_type: DType

Returns the type of the node labels.

node_feature_size property

node_feature_size: int

Returns the dimension of the node labels.

edge_feature_type property

edge_feature_type: DType | None

Returns the type of the edge labels, if present.

edge_feature_size property

edge_feature_size: int | None

Returns the dimension of the edge labels, if present.

matrix_type property

matrix_type: DType

Returns the type of the adjacency matrix.

use_edges property

use_edges: bool

Returns whether the mG compiler expects edge labels in the graph.

tolerance property

tolerance: dict[str, float]

Returns the mapping between types and tolerance values.

use_multiple_loader property

use_multiple_loader: bool

Returns whether the mG compiler expects the usage of the MultipleGraphLoader.

input_spec property

input_spec: tuple[TensorSpec, ...]

Returns the input signature that the mG compiler expects for every model it will produce.

xa_config staticmethod

xa_config(node_config: NodeConfig, matrix_type: tf.DType, tolerance: dict[str, float]) -> CompilerConfig

Returns a CompilationConfig object that tells the mG compiler to expect node labels but no edge labels in the graph and the use of the SingleGraphLoader for the inputs to the model.

Parameters:

• node_config (NodeConfig) –

  The signature of the initial node labels of the graphs.

• matrix_type (DType) –

  The signature of the adjacency matrix of the graphs.

• tolerance (dict[str, float]) –

  The tolerance values for the data types (typically floats) that require an approximate fixed point solution.

Source code in libmg/compiler/compiler.py

@staticmethod
def xa_config(node_config: NodeConfig, matrix_type: tf.DType, tolerance: dict[str, float]) -> CompilerConfig:
    """Returns a ``CompilationConfig`` object that tells the mG compiler to expect node labels but no edge labels in the graph
     and the use of the ``SingleGraphLoader`` for the inputs to the model.

    Args:
        node_config: The signature of the initial node labels of the graphs.
        matrix_type: The signature of the adjacency matrix of the graphs.
        tolerance: The tolerance values for the data types (typically floats) that require an approximate fixed point solution.
    """
    return CompilerConfig(node_config, None, matrix_type, tolerance, False)

xai_config staticmethod

xai_config(node_config: NodeConfig, matrix_type: tf.DType, tolerance: dict[str, float]) -> CompilerConfig

Returns a CompilationConfig object that tells the mG compiler to expect nodes labels but no edge labels in the graph and the use of the MultipleGraphLoader for the inputs to the model.

Parameters:

• node_config (NodeConfig) –

  The signature of the initial node labels of the graphs.

• matrix_type (DType) –

  The signature of the adjacency matrix of the graphs.

• tolerance (dict[str, float]) –

  The tolerance values for the data types (typically floats) that require an approximate fixed point solution.

Source code in libmg/compiler/compiler.py

@staticmethod
def xai_config(node_config: NodeConfig, matrix_type: tf.DType, tolerance: dict[str, float]) -> CompilerConfig:
    """Returns a ``CompilationConfig`` object that tells the mG compiler to expect nodes labels but no edge labels in the graph
     and the use of the ``MultipleGraphLoader`` for the inputs to the model.

    Args:
        node_config: The signature of the initial node labels of the graphs.
        matrix_type: The signature of the adjacency matrix of the graphs.
        tolerance: The tolerance values for the data types (typically floats) that require an approximate fixed point solution.
    """
    return CompilerConfig(node_config, None, matrix_type, tolerance, True)

xae_config staticmethod

xae_config(node_config: NodeConfig, edge_config: EdgeConfig, matrix_type: tf.DType, tolerance: dict[str, float]) -> CompilerConfig

Returns a CompilationConfig object that tells the mG compiler to expect node labels and edge labels in the graph and the use of the SingleGraphLoader for the inputs to the model.

Parameters:

• node_config (NodeConfig) –

  The signature of the initial node labels of the graphs.

• edge_config (EdgeConfig) –

  The signature of the initial edge labels of the graphs.

• matrix_type (DType) –

  The signature of the adjacency matrix of the graphs.

• tolerance (dict[str, float]) –

  The tolerance values for the data types (typically floats) that require an approximate fixed point solution.

Source code in libmg/compiler/compiler.py

@staticmethod
def xae_config(node_config: NodeConfig, edge_config: EdgeConfig, matrix_type: tf.DType, tolerance: dict[str, float]) -> CompilerConfig:
    """Returns a ``CompilationConfig`` object that tells the mG compiler to expect node labels and edge labels in the graph
     and the use of the ``SingleGraphLoader`` for the inputs to the model.

    Args:
        node_config: The signature of the initial node labels of the graphs.
        edge_config: The signature of the initial edge labels of the graphs.
        matrix_type: The signature of the adjacency matrix of the graphs.
        tolerance: The tolerance values for the data types (typically floats) that require an approximate fixed point solution.
    """
    return CompilerConfig(node_config, edge_config, matrix_type, tolerance, False)

xaei_config staticmethod

xaei_config(node_config: NodeConfig, edge_config: EdgeConfig, matrix_type: tf.DType, tolerance: dict[str, float]) -> CompilerConfig

Returns a CompilationConfig object that tells the mG compiler to expect node labels and edge labels in the graph and the use of the MultipleGraphLoader for the inputs to the model.

Parameters:

• node_config (NodeConfig) –

  The signature of the initial node labels of the graphs.

• edge_config (EdgeConfig) –

  The signature of the initial edge labels of the graphs.

• matrix_type (DType) –

  The signature of the adjacency matrix of the graphs.

• tolerance (dict[str, float]) –

  The tolerance values for the data types (typically floats) that require an approximate fixed point solution.

Source code in libmg/compiler/compiler.py

@staticmethod
def xaei_config(node_config: NodeConfig, edge_config: EdgeConfig, matrix_type: tf.DType, tolerance: dict[str, float]) -> CompilerConfig:
    """Returns a ``CompilationConfig`` object that tells the mG compiler to expect node labels and edge labels in the graph
     and the use of the ``MultipleGraphLoader`` for the inputs to the model.

    Args:
        node_config: The signature of the initial node labels of the graphs.
        edge_config: The signature of the initial edge labels of the graphs.
        matrix_type: The signature of the adjacency matrix of the graphs.
        tolerance: The tolerance values for the data types (typically floats) that require an approximate fixed point solution.
    """
    return CompilerConfig(node_config, edge_config, matrix_type, tolerance, True)

MGCompiler

MGCompiler(
    psi_functions: dict[str, Psi | Callable[[], Psi] | Callable[[str], Psi]],
    sigma_functions: dict[str, Sigma | Callable[[], Sigma] | Callable[[str], Sigma]],
    phi_functions: dict[str, Phi | Callable[[], Phi] | Callable[[str], Phi]],
    config: CompilerConfig,
)

The compiler for mG programs.

A program is transformed into a TensorFlow model using the compile method.

Attributes:

• config –

  The configuration for this compiler instance.

• model_inputs –

  The input layers for the models generated by the compiler.

• model_input_spec –

  The input signature for the models generated by the compiler.

• dummy_dataset –

  The dataset used to trace the models generated by the compiler.

• visitor –

  The visitor that traverses an expression tree to construct the model.

Parameters:

• psi_functions (dict[str, Psi | Callable[[], Psi] | Callable[[str], Psi]]) –

  The psi functions that this compiler will recognize.

• sigma_functions (dict[str, Sigma | Callable[[], Sigma] | Callable[[str], Sigma]]) –

  The sigma functions that this compiler will recognize.

• phi_functions (dict[str, Phi | Callable[[], Phi] | Callable[[str], Phi]]) –

  The phi functions that this compiler will recognize.

• config (CompilerConfig) –

  The configuration for the compiler.

Source code in libmg/compiler/compiler.py

def __init__(self, psi_functions: dict[str, Psi | Callable[[], Psi] | Callable[[str], Psi]],
             sigma_functions: dict[str, Sigma | Callable[[], Sigma] | Callable[[str], Sigma]],
             phi_functions: dict[str, Phi | Callable[[], Phi] | Callable[[str], Phi]], config: CompilerConfig):
    """Initializes the instance with the psi, phi and sigma functions that this compiler will recognize and the compiler configuration.

    Args:
        psi_functions: The psi functions that this compiler will recognize.
        sigma_functions: The sigma functions that this compiler will recognize.
        phi_functions: The phi functions that this compiler will recognize.
        config: The configuration for the compiler.
    """
    if config.node_feature_type == tf.float64 or config.edge_feature_type == tf.float64:
        tf.keras.backend.set_floatx('float64')
    elif config.node_feature_type == tf.float16 or config.edge_feature_type == tf.float16:
        tf.keras.backend.set_floatx('float16')
    self.config = config
    self.model_inputs = [tf.keras.Input(shape=(config.node_feature_size, ), name="INPUT_X", dtype=config.node_feature_type),
                         tf.keras.Input(shape=(None,), sparse=True, name="INPUT_A", dtype=config.matrix_type)]
    if config.use_edges:
        self.model_inputs.append(tf.keras.Input(shape=(config.edge_feature_size, ), name="INPUT_E", dtype=config.edge_feature_type))
    if config.use_multiple_loader:
        self.model_inputs.append(tf.keras.Input(shape=(), name="INPUT_I", dtype=tf.int64))
    self.model_input_spec = config.input_spec
    self.dummy_dataset = DummyDataset(NodeConfig(config.node_feature_type, config.node_feature_size), config.matrix_type,
                                      EdgeConfig(config.edge_feature_type, config.edge_feature_size) if config.use_edges else None)  # type: ignore
    self.visitor = self._TreeToTF(FunctionDict(psi_functions), FunctionDict(sigma_functions), FunctionDict(phi_functions), config.tolerance)

compile

compile(expr: str | Tree, verbose: bool = False, memoize: bool = False) -> MGModel

Compiles a mG program into a TensorFlow model.

Parameters:

• expr (str | Tree) –

  The mG program to compile.

• verbose (bool, default: False ) –

  If true, prints some debugging information during the compilation step.

• memoize (bool, default: False ) –

  If true, memoize intermediate outputs during compilation.

Returns:

• MGModel –

  The TensorFlow model that implements expr.

Source code in libmg/compiler/compiler.py

def compile(self, expr: str | Tree, verbose: bool = False, memoize: bool = False) -> MGModel:
    """Compiles a mG program into a TensorFlow model.

    Args:
        expr: The mG program to compile.
        verbose: If true, prints some debugging information during the compilation step.
        memoize: If true, memoize intermediate outputs during compilation.

    Returns:
        The TensorFlow model that implements ``expr``.
    """
    x, a = self.model_inputs[:2]
    e, i = None, None
    if self.config.use_edges and self.config.use_multiple_loader:
        e = self.model_inputs[-2]
        i = self.model_inputs[-1]
    elif self.config.use_edges:
        e = self.model_inputs[-1]
        i = None
    elif self.config.use_multiple_loader:
        e = None
        i = self.model_inputs[-1]

    self.visitor.initialize(IntermediateOutput(mg_parser.parse('__INPUT__'), (x,), a, e, i), memoize)
    tf.keras.backend.clear_session()
    normalized_expr_tree = mg_normalizer.normalize(expr if isinstance(expr, Tree) else mg_parser.parse(expr))
    outputs = self.visitor.visit(normalized_expr_tree)
    model = MGModel(self.model_inputs, outputs.x, normalized_expr_tree, self.visitor.layers, self.config,
                    self.visitor.used_psi, self.visitor.used_phi, self.visitor.used_sigma)
    if verbose is True:
        model.summary(expand_nested=True, show_trainable=True)
    return model

trace

trace(model: MGModel, api: Literal['call', 'predict', 'predict_on_batch']) -> Tuple[MGModel | Callable, float]

Performs tracing on the model, and returns it, together with the time in seconds elapsed for tracing.

The predict_on_batch API cannot be used if the graphs have edge labels as of TF 2.4.

Parameters:

• model (MGModel) –

  The model to trace.

• api (Literal['call', 'predict', 'predict_on_batch']) –

  The TensorFlow API intended to be used with the model. Options are call for the model() API, predict for the model.predict() API and predict_on_batch for the model.predict_on_batch() API.

Returns:

• Tuple[MGModel | Callable, float] –

  The model and the elapsed time in seconds for tracing.

Raises:

• ValueError –

  The predict_on_batch API has been selected and the compiler's configuration is set for graphs with edge labels.

Source code in libmg/compiler/compiler.py

def trace(self, model: MGModel, api: Literal["call", "predict", "predict_on_batch"]) -> Tuple[MGModel | Callable, float]:
    """Performs tracing on the model, and returns it, together with the time in seconds elapsed for tracing.

    The ``predict_on_batch`` API cannot be used if the graphs have edge labels as of TF 2.4.

    Args:
        model: The model to trace.
        api: The TensorFlow API intended to be used with the model. Options are ``call`` for the ``model()`` API,
            ``predict`` for the ``model.predict()`` API and ``predict_on_batch`` for the ``model.predict_on_batch()`` API.

    Returns:
        The model and the elapsed time in seconds for tracing.

    Raises:
        ValueError: The ``predict_on_batch`` API has been selected and the compiler's configuration is set for graphs with edge labels.
    """
    if api == 'predict_on_batch' and self.config.use_edges:
        raise ValueError("The predict_on_batch API, as of TF2.4, isn't compatible with graphs with edge labels.")
    if self.config.use_multiple_loader:
        dummy_loader = MultipleGraphLoader(self.dummy_dataset, batch_size=1, shuffle=False, epochs=1)
    else:
        dummy_loader = SingleGraphLoader(self.dummy_dataset, epochs=1)
    traced_model = MGCompiler._graph_mode_constructor(model, self.model_input_spec, api)
    compile_time = MGCompiler._dummy_run(traced_model, dummy_loader, api)
    return traced_model, compile_time

make_uoperator

make_uoperator(op: Callable[[tf.Tensor], tf.Tensor], name: str | None = None) -> type

Returns a unary operator psi function.

A unary operator is equivalent to a PsiLocal.

Parameters:

• op (Callable[[Tensor], Tensor]) –

  The unary operator function. The function must be compatible with TensorFlow's broadcasting rules. The function is expected to take in input a tensor X with shape (n_nodes, n_node_features), containing the labels of every node in the graph, and return a tensor of shape (n_nodes, n_new_node_features), containing the transformed labels.

• name (str | None, default: None ) –

  The name of the unary operator.

Returns:

• type –

  A subclass of Operator that implements the operator.

Source code in libmg/compiler/functions.py

def make_uoperator(op: Callable[[tf.Tensor], tf.Tensor], name: str | None = None) -> type:
    """Returns a unary operator psi function.

    A unary operator is equivalent to a ``PsiLocal``.

    Args:
        op: The unary operator function. The function must be compatible with TensorFlow's broadcasting
            rules. The function is expected to take in input a tensor X with shape ``(n_nodes, n_node_features)``, containing the labels of
            every node in the graph, and return a tensor of shape ``(n_nodes, n_new_node_features)``, containing the transformed labels.
        name: The name of the unary operator.

    Returns:
        A subclass of Operator that implements the operator.
    """

    def func(self, x: tf.Tensor) -> tf.Tensor:
        return op(x)

    return type(name or "UOperator", (Operator,), {'func': func})

make_boperator

make_boperator(op: Callable[[tf.Tensor, tf.Tensor], tf.Tensor], name: str | None = None) -> type

Returns a binary operator psi function.

A binary operator is a binary local transformation of node labels, psi: (T, T) -> U.

Parameters:

• op (Callable[[Tensor, Tensor], Tensor]) –

  The binary operator function. The function must be compatible with TensorFlow's broadcasting rules. The function is expected to take in input two tensors X1, X2 with shape (n_nodes, n_node_features/2). The first tensor contains for every node the first half of the node features, while the second tensor contains the second half. The operator returns a tensor of shape (n_nodes, n_new_node_features), containing the transformed labels.

• name (str | None, default: None ) –

  The name of the binary operator.

Returns:

• type –

  A subclass of Operator that implements the operator.

Source code in libmg/compiler/functions.py

def make_boperator(op: Callable[[tf.Tensor, tf.Tensor], tf.Tensor], name: str | None = None) -> type:
    """Returns a binary operator psi function.

    A binary operator is a binary local transformation of node labels, psi: (T, T) -> U.

    Args:
        op: The binary operator function. The function must be compatible with TensorFlow's broadcasting
            rules. The function is expected to take in input two tensors X1, X2 with shape ``(n_nodes, n_node_features/2)``. The first tensor contains for every
            node the first half of the node features, while the second tensor contains the second half. The operator returns a tensor of shape
            ``(n_nodes, n_new_node_features)``, containing the transformed labels.
        name: The name of the binary operator.

    Returns:
        A subclass of Operator that implements the operator.
    """

    def func(self, *x: list[tf.Tensor]) -> tf.Tensor:
        return op(x[0], x[1])

    return type(name or "BOperator", (Operator,), {'func': func})

make_koperator

make_koperator(op: Callable[..., tf.Tensor], name: str | None = None) -> type

Returns a k-ary operator psi function.

A k-ary operator is a k-ary local transformation of node labels, psi: (T^k) -> U.

Parameters:

• op (Callable[..., Tensor]) –

  The k-ary operator function. The function must be compatible with TensorFlow's broadcasting rules. The function is expected to take in input k tensors X1, X2, ..., Xk with shape (n_nodes, n_node_features/k). The first tensor contains for every node the first k of the node features, the second tensor contains the next k features, and so on. The operator returns a tensor of shape (n_nodes, n_new_node_features), containing the transformed labels.

• name (str | None, default: None ) –

  The name of the k-ary operator.

Returns:

• type –

  A subclass of Operator that implements the operator.

Source code in libmg/compiler/functions.py

def make_koperator(op: Callable[..., tf.Tensor], name: str | None = None) -> type:
    """Returns a k-ary operator psi function.

    A k-ary operator is a k-ary local transformation of node labels, psi: (T^k) -> U.

    Args:
        op: The k-ary operator function. The function must be compatible with TensorFlow's broadcasting
            rules. The function is expected to take in input k tensors X1, X2, ..., Xk with shape ``(n_nodes, n_node_features/k)``. The first tensor contains
            for every node the first k of the node features, the second tensor contains the next k features, and so on. The operator returns a tensor of shape
            ``(n_nodes, n_new_node_features)``, containing the transformed labels.
        name: The name of the k-ary operator.

    Returns:
        A subclass of Operator that implements the operator.
    """

    def func(self, *x: tf.Tensor) -> tf.Tensor:
        return op(*x)

    return type(name or "KOperator", (Operator,), {'func': func})