def grad_argnums_wrapper(all_vjp_builder):
	"""
	A generic autograd helper funciton. Takes a function that
	builds vjps for all arguments, and wraps it to return only required vjps.
	"""
	def build_selected_vjps(argnums, ans, combined_args, kwargs):
		vjp_func = all_vjp_builder(ans, *combined_args, **kwargs)

		def chosen_vjps(g):
			# Return whichever vjps were asked for
			all_vjps = vjp_func(g)
			return [all_vjps[argnum] for argnum in argnums]

		return chosen_vjps

	return build_selected_vjps