(model, compute_obj_kl, compute_obj_kl_with_grad, compute_kl_hvp)
return wrapper
Makes a wrapper for ngstep for classes that implement nn.Model
    Subsamples inputs for fast Hessian-vector products
Makes a wrapper for ngstep for classes that implement nn.Model
    Subsamples inputs for fast Hessian-vector products
Makes a wrapper for ngstep for classes that implement nn . Model Subsamples inputs for fast Hessian - vector products
def make_ngstep_func(model, compute_obj_kl, compute_obj_kl_with_grad, compute_kl_hvp):
    '''
    Makes a wrapper for ngstep for classes that implement nn.Model
    Subsamples inputs for fast Hessian-vector products
    '''
    assert isinstance(model, nn.Model)

    def wrapper(feed, max_kl, damping, subsample_hvp_frac=.1, grad_stop_tol=1e-6, max_cg_iter=10, enable_bt=True):
        assert isinstance(feed, tuple)

        params0 = model.get_params()
        obj0, kl0, objgrad0 = compute_obj_kl_with_grad(*feed)
        gnorm = util.maxnorm(objgrad0)
        assert np.allclose(kl0, 0), 'Initial KL divergence is %.7f, but should be 0' % (kl0,)
        # Terminate early if gradient is too small
        if gnorm < grad_stop_tol:
            return NGStepInfo(obj0, kl0, obj0, kl0, gnorm, 0)

        # Data subsampling for Hessian-vector products
        subsamp_feed = feed if subsample_hvp_frac is None else subsample_feed(feed, subsample_hvp_frac)
        def hvpx0_func(v):
            with model.try_params(params0):
                hvp_args = subsamp_feed + (v,)
                return compute_kl_hvp(*hvp_args)
        # Objective for line search
        def obj_and_kl_func(p):
            with model.try_params(p):
                obj, kl = compute_obj_kl(*feed)
            return -obj, kl
        params1, num_bt_steps = ngstep(
            x0=params0,
            obj0=-obj0,
            objgrad0=-objgrad0,
            obj_and_kl_func=obj_and_kl_func,
            hvpx0_func=hvpx0_func,
            max_kl=max_kl,
            damping=damping,
            max_cg_iter=max_cg_iter,
            enable_bt=enable_bt)
        model.set_params(params1)
        obj1, kl1 = compute_obj_kl(*feed)
        return NGStepInfo(obj0, kl0, obj1, kl1, gnorm, num_bt_steps)

    return wrapper
def make_ngstep_func ( model , compute_obj_kl , compute_obj_kl_with_grad , compute_kl_hvp ) : assert isinstance ( model , nn . Model ) def wrapper ( feed , max_kl , damping , subsample_hvp_frac = .1 , grad_stop_tol = 1e-6 , max_cg_iter = 10 , enable_bt = True ) : assert isinstance ( feed , tuple ) params0 = model . get_params ( ) obj0 , kl0 , objgrad0 = compute_obj_kl_with_grad ( * feed ) gnorm = util . maxnorm ( objgrad0 ) assert np . allclose ( kl0 , 0 ) , 'Initial KL divergence is %.7f, but should be 0' % ( kl0 , ) # Terminate early if gradient is too small if gnorm < grad_stop_tol : return NGStepInfo ( obj0 , kl0 , obj0 , kl0 , gnorm , 0 ) # Data subsampling for Hessian-vector products subsamp_feed = feed if subsample_hvp_frac is None else subsample_feed ( feed , subsample_hvp_frac ) def hvpx0_func ( v ) : with model . try_params ( params0 ) : hvp_args = subsamp_feed + ( v , ) return compute_kl_hvp ( * hvp_args ) # Objective for line search def obj_and_kl_func ( p ) : with model . try_params ( p ) : obj , kl = compute_obj_kl ( * feed ) return - obj , kl params1 , num_bt_steps = ngstep ( x0 = params0 , obj0 = - obj0 , objgrad0 = - objgrad0 , obj_and_kl_func = obj_and_kl_func , hvpx0_func = hvpx0_func , max_kl = max_kl , damping = damping , max_cg_iter = max_cg_iter , enable_bt = enable_bt ) model . set_params ( params1 ) obj1 , kl1 = compute_obj_kl ( * feed ) return NGStepInfo ( obj0 , kl0 , obj1 , kl1 , gnorm , num_bt_steps ) return wrapper