mark xmap as an API boundary; filter experimental.maps from trace…

…backs and source info

jax/experimental/maps.py

@@ -36,6 +36,7 @@
                                _ensure_index_tuple, donation_vector,
                                shaped_abstractify)
 from jax._src import source_info_util
+from jax._src import traceback_util
 from jax._src.config import config
 from jax.errors import JAXTypeError
 from jax.experimental.global_device_array import GlobalDeviceArray, _get_array_mapping
@@ -52,6 +53,9 @@
                            tuple_insert, moveaxis, split_list, wrap_name)
 from jax import lax
 
+source_info_util.register_exclusion(__file__)
+traceback_util.register_exclusion(__file__)
+
 map, unsafe_map = safe_map, map
 zip = safe_zip
 
@@ -667,7 +671,8 @@ def lower(*args):
         computation, in_tree, in_avals, out_tree(), donate_argnums,
         no_kwargs=True)
 
-  fun_mapped = wraps(fun)(decorate_serial(fun_mapped))
+  fun_mapped = wraps(fun)(
+      traceback_util.api_boundary(decorate_serial(fun_mapped)))
   fun_mapped.lower = decorate_serial(lower)
 
   return fun_mapped