Function relax_check_gradients

tests/python/relax/test_op_gradient_numeric.py:35–204

Generate the forward and the gradient module. Then run them and check numeric gradients. Parameters ---------- op_func : Callable The forward operator function. Should be a function in package relax.op. inputs_numpy : List[np.array] The np array inputs for op_func.

(
    op_func: Callable,
    inputs_numpy: list[np.array],
    target: str | tvm.target.Target,
    dev: tvm.runtime.Device,
    tuple_input: bool = False,
    ignore_grads: list[int] = [],
    **kwargs,  # attr for operators
)

def relax_check_gradients(
    op_func: Callable,
    inputs_numpy: list[np.array],
    target: str | tvm.target.Target,
    dev: tvm.runtime.Device,
    tuple_input: bool = False,
    ignore_grads: list[int] = [],
    **kwargs,  # attr for operators
):
    """Generate the forward and the gradient module. Then run them and check numeric gradients.

    Parameters
    ----------
    op_func : Callable
        The forward operator function. Should be a function in package relax.op.

    inputs_numpy : List[np.array]
        The np array inputs for op_func. inputs_numpy will be transformed into TVM Tensor inside
        this function.

        If op_func takes a tuple of tensors as input, you can set tuple_input as True, and pass the
        tuple input (or list) as inputs_numpy. See test_concat().

    target : Union[str, tvm.target.Target]
        The building target.

    dev : tvm.runtime.Device
        The device to deploy the module.

    tuple_input : bool
        Whether the operator accepts a tuple as input. If true, operator will accept exactly one
        tuple of tensors as input; otherwise, operator accept one or more tensors as input. See
        test_concat(). Default: False.

    ignore_grads: List[int]
        Specifies which input we do not need to find gradient.

        Sometimes the input is not differentiable, such as shape, boolean values, positions, etc.
        We can specify the index of these inputs to check the gradient of them is no_grad, and
        prevent computing numeric gradient.

    kwargs : Any
        The keyword arguments for the op_func. Will be passed to op_func directly.
    """

    func_name = "main"

    # Helper functions
    def _numpy_to_sinfo(data):
        if isinstance(data, list):
            return relax.TupleStructInfo([_numpy_to_sinfo(d) for d in data])
        return relax.TensorStructInfo(data.shape, str(data.dtype))

    def _numpy_to_tvm(data):
        if isinstance(data, list):
            return [_numpy_to_tvm(d) for d in data]
        return tvm.runtime.tensor(data)