//  [python]
def check_approx_error(args, model, params, train_loader_noaug, val_loader, device, total=2000):
	if os.path.isfile(args.save_dir + '/approx_errors.npy'):
		results = np.load(args.save_dir + '/approx_errors.npy')
		draw_matrix(args.save_dir, results[:10], 'approx_error_ntk.pdf')
		draw_matrix(args.save_dir, results[10:], 'approx_error_ella_cov.pdf')
		return

	model2 = copy.deepcopy(model)
	x_full, y_full = subsample(train_loader_noaug, args.num_classes,
							   total+300, args.balanced, device)
	x_test, y_test = x_full[total:total+256], y_full[total:total+256]
	x_full, y_full = x_full[:total], y_full[:total]

	jac_full = jac(model, x_full, y_full, args.num_classes, full=True)
	jac_test = jac(model, x_test, y_test, args.num_classes, full=True).view(
		x_test.shape[0], args.num_classes, -1)

	K1 = jac_full @ jac_full.T
	K1_norm = torch.norm(K1, p=2)

	with torch.no_grad():
		prob_full = model(x_full).softmax(-1).cpu()
		Lambda_full = prob_full.diag_embed() - prob_full[:, :, None] * prob_full[:, None, :]

	JLJ = torch.einsum('NCP,NCD,NDQ->PQ', jac_full.view(total, args.num_classes, -1),
		Lambda_full, jac_full.view(total, args.num_classes, -1))
	Sigma1 = JLJ + torch.eye(JLJ.shape[0]) / args.sigma2
	Sigma1 = Sigma1.inverse()
	# Sigma1_norm = torch.norm(Sigma1, p=2)

	# eq 5
	pred1 = jac_test @ Sigma1 @ jac_test.permute(0, 2, 1)

	results = np.zeros((20, 10))
	for i, M in enumerate([4, 8, 16, 32, 64, 128, 256, 512, 1024, 2000]):
		J_x_subsample = jac(model, x_full[:M], y_full[:M],
							args.num_classes, random=args.random)
		mat = J_x_subsample @ J_x_subsample.T
		p, q = psd_safe_eigen(mat)
		for j, K in enumerate([4, 8, 16, 32, 64, 128, 256, 512, 1024, 2000]):
			if K > M:
				continue

			p = p[range(-1, -(K+1), -1)]
			q = q[:, range(-1, -(K+1), -1)]
			tmp = q.div(p.sqrt())
			V = Js.T @ tmp
			dual_params_list = []
			for item in V.T:
				dual_params = {}
				start = 0
				for name, param in params.items():
					dual_params[name] = item[start:start+param.numel()].view_as(param) #.to(param.device)
					start += param.numel()
				dual_params_list.append(dual_params)
			Psi = partial(Psi_raw, model2, params, dual_params_list)

			psi_full = Psi(x_full).cpu()
			psi_test = Psi(x_test).cpu()

			# nystrom approximation of K1
			K2 = psi_full.flatten(0, 1) @ psi_full.flatten(0, 1).T

			# normalized E'
			results[i, j] = torch.norm(K1 - K2, p=2) / K1_norm

			# Eq 8
			G = torch.einsum('NCP,NCD,NDQ->PQ', psi_full, Lambda_full, psi_full) + torch.eye(psi_full.shape[-1]) / args.sigma2
			pred2 = psi_test @ G.inverse() @ psi_test.permute(0, 2, 1)
			results[i+10, j] = (torch.norm(pred1 - pred2, p=2, dim=(-2, -1)) / torch.norm(pred1, p=2, dim=(-2, -1))).mean()

			print(i, j, results[i, j], results[i+10, j])

			if 0 and K == M:
				# the explicit nystrom approximation
				K3 = jac_full @ J_x_subsample.T @ (J_x_subsample @ J_x_subsample.T).pinverse() @ (J_x_subsample @ jac_full.T)
				print('check1', torch.dist(K2, K3), K2[:5,:5], K3[:5,:5])

				# Sigma' in eq 21
				Sigma2 = (J_x_subsample @ JLJ @ J_x_subsample.T + (J_x_subsample @ J_x_subsample.T) / args.sigma2).inverse()
				Sigma2 = J_x_subsample.T @ Sigma2 @ J_x_subsample

				# another form of Sigma' in Appendix A.4
				P = J_x_subsample.T @ (J_x_subsample @ J_x_subsample.T).inverse() @ J_x_subsample
				Sigma3 = (P @ JLJ @ P + torch.eye(JLJ.shape[0]) / args.sigma2).inverse() + args.sigma2 * (P - torch.eye(JLJ.shape[0]))
				print('check2', torch.dist(Sigma2, Sigma3), Sigma2[:5, :5], Sigma3[:5, :5])

				# eq 21 equals to eq 8
				pred3 = jac_test @ Sigma2 @ jac_test.permute(0, 2, 1)
				print('check3', torch.dist(pred2, pred3), pred2[5], pred3[5])

				# E
				E = torch.norm(Sigma1 - Sigma2, p=2)
				print('E', torch.dist(Sigma1, Sigma2), E, K1_norm, E / K1_norm)
				print(Sigma1[:5, :5], Sigma2[:5, :5])

	np.save(args.save_dir + '/approx_errors', results)
	draw_matrix(args.save_dir, results[:10], 'approx_error_ntk.pdf')
	draw_matrix(args.save_dir, results[10:], 'approx_error_ella_cov.pdf')