def _axis_index_bind(*, axis_name):
if not isinstance(axis_name, (tuple, list)):
axis_name = (axis_name,)
inner_size = 1
index = 0
for name in reversed(axis_name):
frame = core.axis_frame(name)
if frame.main_trace is not None:
trace = frame.main_trace.with_cur_sublevel()
name_idx = trace.process_axis_index(frame)
else:
name_idx = core.Primitive.bind(axis_index_p, axis_name=name)
index += name_idx * inner_size
inner_size *= psum(1, name)
return index
def p_func_if_pmap(obj, axis_name):
try:
core.axis_frame(axis_name)
return p_func(obj, axis_name)
except NameError:
return obj