Fix Distrax from Jax API change.

The method map_primitive was removed from Primitive class.

PiperOrigin-RevId: 888552312

Author: DistraxDev

distrax/_src/utils/transformations.py

@@ -331,7 +331,9 @@ def write(var, val):
 
 
 def _extract_call_jaxpr(primitive, params):
-  if not (primitive.call_primitive or primitive.map_primitive):
+  if not (
+      primitive.call_primitive or getattr(primitive, "map_primitive", False)
+  ):
     return None, params
   else:
     params = dict(params)