(self, *args, **kwargs)
| 345 | return xc._xla.buffer_to_dlpack_managed_tensor(x, take_ownership=take_ownership) |
| 346 | |
| 347 | def execute(self, *args, **kwargs): |
| 348 | from ..tensor import Tensor |
| 349 | from ..optimizer import Optimizer |
| 350 | from ..traced_module.pytree import tree_flatten |
| 351 | from ..utils.module_utils import get_expand_structure |
| 352 | |
| 353 | inputs, _ = tree_flatten((args, kwargs)) |
| 354 | arrays = [] |
| 355 | cn = CompNode(get_default_device()) |
| 356 | stream = dict(self.xla_exec.backend.get_compute_compnode()) |
| 357 | device_kind, device_id, stream_id = cn.physical_locator |
| 358 | |
| 359 | xla_stream = stream[device_id] |
| 360 | xla_comp_cn = "gpu{}:{}".format(device_id, xla_stream) |
| 361 | self.optimizers = [] |
| 362 | for t in inputs: |
| 363 | if isinstance(t, RawTensor): |
| 364 | if not t._is_external_value(): |
| 365 | assert cn == t.device |
| 366 | arrays.append(t.to(xla_comp_cn, _borrow=True)) |
| 367 | else: |
| 368 | arrays.append(t) |
| 369 | if isinstance(t, Optimizer): |
| 370 | self.optimizers.append(t) |
| 371 | |
| 372 | arrays = self.prepare_xla_inputs(arrays) |
| 373 | outputs = self.xla_exec(*arrays) |
| 374 | global xla_client_compute_stream |
| 375 | xla_client_compute_stream = xla_stream |
| 376 | return_vals = [] |
| 377 | for i in self.out_list: |
| 378 | if i == -1: |
| 379 | if not hasattr(self, "outdef"): |
| 380 | return_vals.append(None) |
| 381 | else: |
| 382 | return_vals.append(outputs[self.outkey2idx[i]]) |
| 383 | if not self.out_list: |
| 384 | return_vals = [ |
| 385 | None, |
| 386 | ] |
| 387 | keeped_features = [] |
| 388 | for i in self.keeped_activation: |
| 389 | keeped_features.append(tensor(outputs[self.outkey2idx[i]], device=cn)) |
| 390 | out_tensors = [] |
| 391 | for array in return_vals: |
| 392 | if array is not None: |
| 393 | t = tensor(array, device=cn) |
| 394 | out_tensors.append(t) |
| 395 | else: |
| 396 | out_tensors.append(array) |
| 397 | if self.overall: |
| 398 | for attr, key in self.update_param_dict.items(): |
| 399 | param = get_expand_structure(attr[0], attr[1]) |
| 400 | xla_array = outputs[self.outkey2idx[key]] |
| 401 | t = tensor(xla_array, device=cn) |
| 402 | param._reset(t) |
| 403 | |
| 404 | for state, key in self.update_opt_param_dict.items(): |
nothing calls this directly
no test coverage detected