Method Python

caffe2/python/core.py:2296–2374

Registers and returns a python operator. `f` and `grad_f` can be one of the following: - a function with signature (inputs, outputs), where inputs and outputs are a list of CPUTensor objects. This function will be called from C++ everytime th

(
        self,
        f,
        grad_f=None,
        python_func_type=None,
        pass_workspace=False,
        grad_output_indices=None,
        grad_input_indices=None
    )

        )))

    def Python(
        self,
        f,
        grad_f=None,
        python_func_type=None,
        pass_workspace=False,
        grad_output_indices=None,
        grad_input_indices=None
    ):
        """
        Registers and returns a python operator.

        `f` and `grad_f` can be one of the following:
            - a function with signature (inputs, outputs), where inputs and
              outputs are a list of CPUTensor objects. This function will be
              called from C++ everytime the operator is executed.
            - a tuple (func, args, kwargs), here `func` is a callable, args is
              an argument list, and kwargs is a dict list. The call:
                  f = func(*args, kwargs)
              will be performed locally at node initialization time, on all of
              the nodes of the job, returning `f`, a callable that will be used
              as the python operator function to be called during Net execution.
              This is to be used when using python operator in a distributed
              context, and allows to create and keep local python state across
              calls to the operator.

        `python_func_type` is a type of an object that constructed as
        python_func_type(f) and provides an implementation to forward and
        backward functions. Its useful in such a case where users needs
        a statefull PythonOp (ex: use autograd for computing grad_f).

        If `pass_workspace` is True, the signature is changed to
        (inputs, outputs, workspace) where `workspace` is the workspace the op
        is going to run on. This is potentially dangerous (as the op can
        manipulate the workspace directly), use on your own risk.

        If a gradient function is specified (`grad_f`), by default its inputs
        will be: (1) all inputs to `f`, (2) followed by all outputs of `f`, (3)
        and then all gradient outputs of `f`. The outputs of `grad_f` will be
        (by default) all gradient inputs to `f`. If a subset of the gradient
        outputs or gradient inputs is desired instead, then the subsets can be
        specified by providing `grad_output_indices` and/or `grad_input_indices`
        which identify the indices of `f`'s inputs and outputs which have
        gradients.
        """
        assert(IsOperator('Python'))

        def make_builder(t):
            if not isinstance(t, tuple):
                return ''
            assert len(t) == 3, 'Expected builder tuple (func, args, kwargs)'
            func, args, kwargs = t
            normalized = (func, tuple(args), dict(kwargs))
            return pickle.dumps(normalized)

        f_builder = make_builder(f)
        grad_f_builder = make_builder(grad_f)