def _check_set(x):
x, i = x
# TODO (bozhi) need C pointers!
def _hard_set(i, value):
if len(i) == 1: self._latest_view[i[0]] = value
elif len(i) == 2: self._latest_view[i[0]][i[1]] = value
else:
raise NotImplementedError(
"High-dimensional PartitionedTensor with None not supported!"
)
if x is None:
_hard_set(i, value)
return
is_to_array = is_array(x)
is_from_array = is_array(value)
if is_from_array and is_to_array:
try:
copy(x, value)
except ValueError:
warn(
"Incompatible arrays (e.g. different shapes). Overwritting."
)
_hard_set(i, value)
else:
# TODO (bozhi): should not allow None assignment but implement free(index)
if not is_to_array and x is not None:
warn("Array partition was modified as %s object." %
type(x))
if not is_from_array and value is not None:
warn("Modifying array partition with %s object!" %
type(value))
x = value
def get_placement_for_any(placement: Union[Collection[Union[Architecture, Device, Task, "TaskID", Any]], Any, None]) \
-> FrozenSet[Device]:
if placement is not None:
ps = placement if isinstance(placement, Iterable) and not array.is_array(placement) else [placement]
return get_placement_for_set(ps)
else:
return frozenset(get_all_devices())
def get_placement_for_value(p: Union[Architecture, Device, Task, TaskID, Any]) -> List[Device]:
if hasattr(p, "__parla_placement__"):
# this handles Architecture, ResourceRequirements, and other types with __parla_placement__
return list(p.__parla_placement__())
elif isinstance(p, Device):
return [p]
elif isinstance(p, TaskID):
return get_placement_for_value(p.task)
elif isinstance(p, task_runtime.Task):
return get_placement_for_value(p.req)
elif array.is_array(p):
return [array.get_memory(p).device]
elif isinstance(p, Collection):
raise TypeError("Collection passed to get_placement_for_value, probably needed get_placement_for_set: {}"
.format(type(p)))
else:
raise TypeError(type(p))
def __getitem__(self, index: IndexType): # -> Union[Array, List[Array]]
"""
Read partitions and make sure they are on the current device.
:param index: index of the target partition(s).
.. todo:
Multiple partitions are currently returned as a Python list of partitions (ndarrays).
"""
if not isinstance(index, tuple):
index = (index, )
ret = []
parse_index(
self._latest_view,
index,
step=lambda I, i: I[i],
stop=lambda x: ret.append(clone_here(x) if is_array(x) else x))
if len(ret) == 1:
if ret[0] is None: warn("Partition has been freed!")
return ret[0]
warn(
"Multiple partitions are currently returned as a Python list of partitions (ndarrays)."
)
return ret