code: def size(a, axis=None):
      ivy.set_default_int_dtype("int64")
      if axis is not None:
          sh = ivy.shape(a)
          return sh[axis]
      return a.size
  
imports: import ivy
  from ivy.func_wrapper import with_unsupported_dtypes
  from ivy.functional.frontends.jax.array import Array
  import ivy.functional.frontends.jax.numpy
  from ivy.functional.frontends.jax.func_wrapper import to_ivy_arrays_and_back, outputs_to_frontend_arrays, handle_jax_dtype, inputs_to_ivy_arrays
  from ivy.func_wrapper import handle_out_argument
  from ivy import with_unsupported_device_and_dtypes
label_desc: Use image node_id 30 for calling a global function with example usage: size(a, axis) and returns: a, sh