Context-manager forcing placement of ops and Tensors on a device.
| 1482 | |
| 1483 | |
| 1484 | class _EagerDeviceContext(object): |
| 1485 | """Context-manager forcing placement of ops and Tensors on a device.""" |
| 1486 | |
| 1487 | def __init__(self, ctx, device_name): |
| 1488 | self._device_name = device_name |
| 1489 | self._ctx = ctx |
| 1490 | self._stack = [] |
| 1491 | |
| 1492 | def __enter__(self): |
| 1493 | ctx = self._ctx |
| 1494 | old_device_name = ctx.device_name |
| 1495 | old_device_spec = ctx.device_spec |
| 1496 | new_device_name = self._device_name |
| 1497 | cache_key = (old_device_name, new_device_name) |
| 1498 | try: |
| 1499 | new_device_name, new_device_spec = _device_parsing_cache[cache_key] |
| 1500 | except TypeError: |
| 1501 | # Error while trying to compute the cache key. |
| 1502 | raise ValueError("Expecting a string device name. Got %s(%s)" % |
| 1503 | (type(new_device_name), new_device_name)) |
| 1504 | except KeyError: |
| 1505 | # Handle a cache miss. |
| 1506 | if new_device_name is not None: |
| 1507 | if not isinstance(new_device_name, six.string_types): |
| 1508 | raise ValueError("Expecting a string device name. Got %s(%s)" % |
| 1509 | (type(new_device_name), new_device_name)) |
| 1510 | device_spec = pydev.DeviceSpec.from_string(new_device_name) |
| 1511 | if old_device_name: |
| 1512 | new_device_spec = copy.copy(old_device_spec) |
| 1513 | else: |
| 1514 | ctx.ensure_initialized() |
| 1515 | new_device_spec = pydev.DeviceSpec.from_string( |
| 1516 | ctx._context_devices[0]) # pylint: disable=protected-access |
| 1517 | new_device_spec = new_device_spec.make_merged_spec(device_spec) |
| 1518 | else: |
| 1519 | new_device_spec = pydev.DeviceSpec.from_string("") |
| 1520 | new_device_name = new_device_spec.to_string() |
| 1521 | _device_parsing_cache[cache_key] = (new_device_name, new_device_spec) |
| 1522 | |
| 1523 | ctx._set_device(new_device_name, new_device_spec) # pylint: disable=protected-access |
| 1524 | self._stack.append((old_device_name, old_device_spec, new_device_spec)) |
| 1525 | |
| 1526 | def __exit__(self, *ex_info): |
| 1527 | ctx = self._ctx |
| 1528 | old_device_name, old_device_spec, new_device_spec = self._stack[-1] |
| 1529 | if ctx.device_spec is not new_device_spec: |
| 1530 | raise RuntimeError( |
| 1531 | "Exiting device scope without proper scope nesting") |
| 1532 | del self._stack[-1] |
| 1533 | ctx._set_device(old_device_name, old_device_spec) # pylint: disable=protected-access |
| 1534 | |
| 1535 | |
| 1536 | # Do not set directly. Use _set_context. |