def METHOD_NAME(op: tf.Operation):
    # Make sure to not wrap it multiple times.
    if getattr(op, "_RETURNN_gradient_checkpoint_wrapped_grad_func", None) is not None:
        return
    try:
        orig_grad_func = ops.get_gradient_function(op)
    except LookupError:
        return  # no gradient function
    if orig_grad_func is None:
        return
    class _WrappedOp:
        def __init__(self):
            self.op = op
            self._inputs = tuple(_map_tensor(x) for x in op.inputs)
            self._outputs = tuple(_map_tensor(x) for x in op.outputs)
        @property
        def inputs(self) -> Tuple[tf.Tensor, ...]:
            """inputs"""
            return self._inputs
        @property
        def outputs(self) -> Tuple[tf.Tensor, ...]:
            """outputs"""
            return self._outputs
        def get_attr(self, name: str):
            """get_attr"""
            return self.op.get_attr(name)
    _wrapped_op = _WrappedOp()
    # noinspection PyShadowingNames
    def _wrapped_grad_func(op, *out_grads):
        assert op is _wrapped_op.op  # wrapped grad func not intended to be called for other ops
        return orig_grad_func(_wrapped_op, *out_grads)
    op._gradient_function = _wrapped_grad_func
    op._RETURNN_gradient_checkpoint_wrapped_grad_func = _wrapped_grad_func