def get_aggregate_fn(agg):
    """Internal function to get the aggregation function for node data
    generated from different relations.

    Parameters
    ----------
    agg : str
        Method for aggregating node features generated by different relations.
        Allowed values are 'sum', 'max', 'min', 'mean', 'stack'.

    Returns
    -------
    callable
        Aggregator function that takes a list of tensors to aggregate
        and returns one aggregated tensor.
    """
    if agg == "sum":
        fn = nd.sum
    elif agg == "max":
        fn = nd.max
    elif agg == "min":
        fn = nd.min
    elif agg == "mean":
        fn = nd.mean
    elif agg == "stack":
        fn = None  # will not be called
    else:
        raise DGLError(
            "Invalid cross type aggregator. Must be one of "
            '"sum", "max", "min", "mean" or "stack". But got "%s"' % agg
        )
    if agg == "stack":

        def stack_agg(inputs, dsttype):  # pylint: disable=unused-argument
            if len(inputs) == 0:
                return None
            return nd.stack(*inputs, axis=1)

        return stack_agg
    else:

        def aggfn(inputs, dsttype):  # pylint: disable=unused-argument
            if len(inputs) == 0:
                return None
            stacked = nd.stack(*inputs, axis=0)
            return fn(stacked, axis=0)

        return aggfn
Internal function to get the aggregation function for node data
    generated from different relations.

    Parameters
    ----------
    agg : str
        Method for aggregating node features generated by different relations.
        Allowed values are 'sum', 'max', 'min', 'mean', 'stack'.

    Returns
    -------
    callable
        Aggregator function that takes a list of tensors to aggregate
        and returns one aggregated tensor.
    
python/dgl/nn/mxnet/hetero.py