From d3a378f9665d8eee827c74cb9ecbee81e4c288dd Mon Sep 17 00:00:00 2001 From: Bruce Fontaine Date: Wed, 30 Sep 2020 12:40:59 -0700 Subject: [PATCH] Add EPU device type. Allow passing string to initialize_system_for_tpu_embedding. PiperOrigin-RevId: 334657063 Change-Id: I6c339b32fb118fd9d8ac93c55656fe6915484aae --- tensorflow/python/framework/device_spec.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/framework/device_spec.py b/tensorflow/python/framework/device_spec.py index 36d30450e92..8b7d72bf03a 100644 --- a/tensorflow/python/framework/device_spec.py +++ b/tensorflow/python/framework/device_spec.py @@ -21,7 +21,8 @@ from __future__ import print_function from tensorflow.python.util.tf_export import tf_export -_VALID_DEVICE_TYPES = frozenset({"CPU", "GPU", "TPU", "CUSTOM"}) +# EPU represents for TPU embedding for now. Subject to change in future. +_VALID_DEVICE_TYPES = frozenset({"CPU", "GPU", "TPU", "CUSTOM", "EPU"}) # ==============================================================================