code: def average(
      a, axis=None, weights=None, returned=False, keepdims=False
  ):
      # canonicalize_axis to ensure axis or the values in axis > 0
      if isinstance(axis, (tuple, list)):
          a_ndim = len(ivy.shape(a))
          new_axis = [0] * len(axis)
          for i, v in enumerate(axis):
              if not -a_ndim <= v < a_ndim:
                  raise ValueError(
                      f"axis {v} is out of bounds for array of dimension {a_ndim}"
                  )
              new_axis[i] = v + a_ndim if v < 0 else v
          axis = tuple(new_axis)
  
      if weights is None:
          ret = ivy.mean(a, axis=axis, keepdims=keepdims)
          if axis is None:
              fill_value = (
                  int(a.size)
                  if ivy.is_int_dtype(ret)
                  else float(a.size)
              )
              weights_sum = ivy.full((), fill_value, dtype=ret.dtype)
          else:
              if isinstance(axis, tuple):
                  # prod with axis has dtype Sequence[int]
                  fill_value = 1
                  for d in axis:
                      fill_value *= a.shape[d]
              else:
                  fill_value = a.shape[axis]
              weights_sum = ivy.full_like(ret, fill_value=fill_value)
      else:
          a = ivy.asarray(a, copy=False)
          weights = ivy.asarray(weights, copy=False)
          a, weights = promote_types_of_jax_inputs(a, weights)
  
          a_shape = ivy.shape(a)
          a_ndim = len(a_shape)
          weights_shape = ivy.shape(weights)
  
          # Make sure the dimensions work out
          if a_shape != weights_shape:
              if len(weights_shape) != 1:
                  raise ValueError(
                      "1D weights expected when shapes of a and weights differ."
                  )
              if axis is None:
                  raise ValueError(
                      "Axis must be specified when shapes of a and weights differ."
                  )
              elif isinstance(axis, tuple):
                  raise ValueError(
                      "Single axis expected when shapes of a and weights differ"
                  )
              elif weights.shape[0] != a.shape[axis]:
                  raise ValueError(
                      "Length of weights not compatible with specified axis."
                  )
  
              weights = ivy.broadcast_to(
                  weights, shape=(a_ndim - 1) * (1,) + weights_shape
              )
              weights = ivy.moveaxis(weights, -1, axis)
  
          weights_sum = ivy.sum(weights, axis=axis)
          ret = (
              ivy.sum(a * weights, axis=axis, keepdims=keepdims)
              / weights_sum
          )
  
      if returned:
          if ret.shape != weights_sum.shape:
              weights_sum = ivy.broadcast_to(
                  weights_sum, shape=ret.shape
              )
          return ret, weights_sum
  
      return ret
  
imports: import ivy
  from ivy.func_wrapper import with_unsupported_dtypes
  from ivy.functional.frontends.jax.func_wrapper import to_ivy_arrays_and_back, handle_jax_dtype
  from ivy.functional.frontends.jax.numpy import promote_types_of_jax_inputs
  from ivy.functional.backends.jax.experimental.elementwise import _normalize_axis_tuple
label_desc: Use image node_id 2 for calling a global function with example usage: average(a, axis, weights, returned, keepdims) and returns: ret, ret,  weights_sum