code: def _forward_fn(
      xs,
      x,
      func,
      duplicate_index_chains,
      xs_grad_idxs=None,
      ret_grad_idxs=None,
  ):
      """Forward function for gradient calculation."""
      # Setting x(relevant variables) into xs(all variables)
      x = ivy.nested_map(ivy.to_ivy, x, include_derived=True)
      x_arr_idxs = ivy.nested_argwhere(x, ivy.is_array)
      x_arr_values = ivy.multi_index_nest(x, x_arr_idxs)
      if xs_grad_idxs is not None:
          xs_grad_arr_idxs = []
          for grad_idx in xs_grad_idxs:
              xs_grad_arr_idx = ivy.nested_argwhere(
                  ivy.index_nest(xs, grad_idx), ivy.is_array
              )
              for idx in xs_grad_arr_idx:
                  xs_grad_arr_idxs.append(list(grad_idx) + idx)
          ivy.set_nest_at_indices(xs, xs_grad_arr_idxs, x_arr_values)
      elif ivy.is_array(xs):
          xs = x
      else:
          xs_arr_idxs = ivy.nested_argwhere(
              xs, lambda x: ivy.is_array(x)
          )
          ivy.set_nest_at_indices(xs, xs_arr_idxs, x_arr_values)
  
      # Setting duplicates to ensure same references as in the original input
      if not ivy.is_array(xs):
          xs = _set_duplicates(xs, duplicate_index_chains)
      ret = func(xs)
  
      # Getting the relevant outputs from the function return for gradient calculation
      _, ret_values = _get_native_variables_and_indices(
          ret, idxs=ret_grad_idxs
      )
      if (
          isinstance(ret_values, list)
          and len(ret_values) == 1
          and ret_grad_idxs is None
      ):
          ret_values = ret_values[0]
      return ret_values
  
file_path: ivy/ivy/functional/backends/jax/gradients.py
imports: import jax
  import jax.lax
  from ivy.functional.backends.jax import JaxArray, NativeArray
  from typing import Optional, Callable, Sequence, Union, Tuple
  import ivy
  from ivy.functional.ivy.gradients import _get_required_float_variables, _get_y_and_ret_idxs, _get_native_variables_and_indices, _set_duplicates, _process_func_ret_and_grads
label_desc: Use image node_id 4 for calling a global function with example usage: _forward_fn(xs, x, func, duplicate_index_chains, xs_grad_idxs, ret_grad_idxs) and returns: ret_values