human
I wrote the below code

    layer_scopes = []
    def rev_block_grad(op, grad_y):
        'Custom    gradient    fn    for    a    block    of    reversible    residual    layers.'
        y = op.outputs[0]
        ys = tf.split(y, 2, axis=(len(y.get_shape()) - 1))
        grad_ys = tf.split(grad_y, 2, axis=(len(y.get_shape()) - 1))
        all_vars = op.inputs[1:]
        f_vars = [[] for _ in range(num_layers)]
        g_vars = [[] for _ in range(num_layers)]
        f_vars_idxs = [[] for _ in range(num_layers)]
        g_vars_idxs = [[] for _ in range(num_layers)]
        for (i, v) in enumerate(all_vars):
            ref = v.op.inputs[0]
            assert (ref.dtype == dtypes.float32_ref)
            regex = LAYER_RE.match(v.name)
            layer_no = int(regex.group(1))
            fn_name = regex.group(2)
            if (fn_name == 'f'):
                f_vars[layer_no].append(ref)
                f_vars_idxs[layer_no].append(i)
            else:
                assert (fn_name == 'g')
                g_vars[layer_no].append(ref)
                g_vars_idxs[layer_no].append(i)
        f_grads = []
        g_grads = []
        layer_scopes.reverse()
        f_vars.reverse()
        g_vars.reverse()
        for i in xrange(num_layers):
            with tf.variable_scope(layer_scopes[i], reuse=True):
                (ys, grad_ys, grad_f_vars, grad_g_vars) = _rev_layer_backward(ys, grad_ys, f, g, f_vars[i], g_vars[i])
                f_grads.append(grad_f_vars)
                g_grads.append(grad_g_vars)
        f_grads.reverse()
        g_grads.reverse()
        var_grads = ([None] * len(all_vars))
        for (idxs, grads) in (zip(f_vars_idxs, f_grads) + zip(g_vars_idxs, g_grads)):
            for (i, grad) in zip(idxs, grads):
                var_grads[i] = grad
        grad_x = tf.concat(grad_ys, axis=(-1))
        all_grads = ([grad_x] + var_grads)
        return all_grads
    @function.Defun(tf.float32, python_grad_func=rev_block_grad, shape_func=(lambda _: [x.get_shape()]))
    def rev_block_defun(inp):
       ...

What's a good function header?
 ... Here's a good function header?
assistant
#Here's a good function header.

def rev_block(x, f, g, num_layers=1, is_training=True):