code: def _custom_vjp_single_scope_fn(
      fn: Callable[..., Any],
      backward_fn: Callable[..., Any],
      grad_vars: lift.CollectionFilter = "params",
      nondiff_argnums=(),
  ):
      nodiff_fn = functools.partial(fn, needs_residual=False)
      forward_fn = functools.partial(fn, needs_residual=True)
      return lift.custom_vjp(
          nodiff_fn, forward_fn, backward_fn, grad_vars, nondiff_argnums
      )
  
imports: import dataclasses
  import functools
  import inspect
  from typing import Any, Callable, Dict, Iterable, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union
  from flax import errors, struct, traceback_util
  from flax.core import Scope, lift, meta
  from flax.core.frozen_dict import FrozenDict
  from flax.linen import module
  from flax.linen.module import Module, Variable, _derive_profiling_name, _get_unbound_fn, wrap_method_once
  import jax
label_desc: Use image node_id 26 for calling a global function with example usage: _custom_vjp_single_scope_fn(fn, backward_fn, grad_vars, nondiff_argnums) and returns: lift