code: def fwd_grad_chooser(g, ans, x, axis=None, keepdims=False):
      if anp.isscalar(x):
          return g
      if not keepdims:
          if isinstance(axis, int):
              ans = anp.expand_dims(ans, axis)
          elif isinstance(axis, tuple):
              for ax in sorted(axis):
                  ans = anp.expand_dims(ans, ax)
      chosen_locations = x == ans
      return anp.sum(
          (g * chosen_locations), axis=axis, keepdims=keepdims
      ) / anp.sum(chosen_locations, axis=axis, keepdims=keepdims)
  
file_path: autograd/autograd/numpy/numpy_jvps.py
imports: from .None import numpy_wrapper
  from .numpy_vjps import untake, balanced_eq, match_complex, replace_zero, dot_adjoint_0, dot_adjoint_1, tensordot_adjoint_0, tensordot_adjoint_1, nograd_functions
  from autograd.extend import defjvp, defjvp_argnum, def_linear, vspace, JVPNode, register_notrace
  from ..util import func
  from .numpy_boxes import ArrayBox
label_desc: Use image node_id 3 for calling a global function with example usage: fwd_grad_chooser(g, ans, x, axis, keepdims) and returns: unknown, g