Function register_func

python/dgl/_ffi/function.py:147–206

Register global function Parameters ---------- func_name : str or function The function name f : function, optional The function to be registered. override: boolean optional Whether override existing entry. Returns ------- fregister : funct

(func_name, f=None, override=False)

def register_func(func_name, f=None, override=False):
    """Register global function

    Parameters
    ----------
    func_name : str or function
        The function name

    f : function, optional
        The function to be registered.

    override: boolean optional
        Whether override existing entry.

    Returns
    -------
    fregister : function
        Register function if f is not specified.

    Examples
    --------
    The following code registers my_packed_func as global function.
    Note that we simply get it back from global function table to invoke
    it from python side. However, we can also invoke the same function
    from C++ backend, or in the compiled DGL code.

    .. code-block:: python

      targs = (10, 10.0, "hello")
      @dgl.register_func
      def my_packed_func(*args):
          assert(tuple(args) == targs)
          return 10
      # Get it out from global function table
      f = dgl.get_global_func("my_packed_func")
      assert isinstance(f, dgl.nd.Function)
      y = f(*targs)
      assert y == 10
    """
    if callable(func_name):
        f = func_name
        func_name = f.__name__

    if not isinstance(func_name, str):
        raise ValueError("expect string function name")

    ioverride = ctypes.c_int(override)

    def register(myf):
        """internal register function"""
        if not isinstance(myf, Function):
            myf = convert_to_dgl_func(myf)
        check_call(
            _LIB.DGLFuncRegisterGlobal(c_str(func_name), myf.handle, ioverride)
        )
        return myf

    if f: