def thrust(obj):
if isinstance(obj, list):
result = [thrust(o) for o in obj]
elif hasattr(obj, 'data'):
result = obj.data
elif isinstance(obj, float):
result = PrecisionResolver.get_floating_point(obj)
elif isinstance(obj, int):
result = trtc.DVInt64(obj)
else:
raise ValueError(f"Cannot upload {obj} to device.")
return result
def __setitem__(self, key, value):
if hasattr(value, 'data') and hasattr(value, 'shape') and len(value.shape) != 0:
if isinstance(value, np.ndarray):
vector = trtc.device_vector_from_numpy(value)
trtc.Copy(vector, self.data)
else:
trtc.Copy(value.data, self.data)
else:
if isinstance(value, int):
dvalue = trtc.DVInt64(value)
elif isinstance(value, float):
dvalue = PrecisionResolver.get_floating_point(value)
else:
raise TypeError("Only Storage, int and float are supported.")
trtc.Fill(self.data, dvalue)
return self