def set_parallel_axes(axes, parallel_axis):
new_axes = []
for axis in Axes.as_nested_list(axes):
if axis == parallel_axis:
axis = parallel_axis
elif isinstance(axis, collections.Iterable):
# flattened axis
axis = [parallel_axis if a == parallel_axis else a for a in axis]
new_axes.append(axis)
return make_axes(new_axes)
def set_parallel_axes(axes, parallel_axis):
new_axes = []
flat_names = dict()
for i, axis in enumerate(Axes.as_nested_list(axes)):
if axis == parallel_axis:
axis = parallel_axis
elif isinstance(axis, collections.Iterable):
flat_names[i] = axes[i].name
axis = [parallel_axis if a == parallel_axis else a for a in axis]
new_axes.append(axis)
new_axes = make_axes(new_axes)
for i in flat_names:
new_axes[i].name = flat_names[i]
return new_axes
def generate_default_lut_layout(op):
"""
Generates the default layout assignment for a lookup table operation on GPU.
Arguments:
op (LookupTableOp): op to generate layout for
Returns:
GPULayoutAssignment for this op
"""
axes_list = Axes.as_flattened_list(op.axes)
groups = Axes.as_nested_list(op.axes)
layout = []
for group in groups:
if isinstance(group, list):
layout.append([axes_list.index(a) for a in group])
else:
layout.append([axes_list.index(group)])
return [GPULayoutAssignment(axes_list, layout)]