tf-models-nightly 2.17.0.dev20240624__py2.py3-none-any.whl → 2.17.0.dev20240625__py2.py3-none-any.whl

official/common/distribute_utils.py

@@ -148,8 +148,32 @@ def get_distribution_strategy(distribution_strategy="mirrored",
 
   if distribution_strategy == "tpu":
     # When tpu_address is an empty string, we communicate with local TPUs.
-    cluster_resolver = tpu_initialize(tpu_address)
-    return tf.distribute.TPUStrategy(cluster_resolver)
+    # Bug workaround that in v5p we need to explicitly specify the device
+    # assignment when using tpu strategy, adding device assignment to the
+    # strategy.
+    cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
+        tpu=tpu_address
+    )
+    if tpu_address not in ("", "local"):
+      tf.config.experimental_connect_to_cluster(cluster_resolver)
+    topology = tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
+
+    device_assignment = None
+    if hasattr(tf.tpu.experimental, "HardWareFeature"):
+      hardware_feature = tf.tpu.experimental.HardWareFeature(
+          cluster_resolver.tpu_hardware_feature
+      )
+      if (
+          hardware_feature.embedding_feature
+          == tf.tpu.experimental.HardwareFeature.EmbeddingFeature.V2
+      ):
+        tpu_metadata = cluster_resolver.get_tpu_system_metadata()
+        device_assignment = tf.tpu.experimental.DeviceAssignment.build(
+            topology, num_replicas=tpu_metadata.num_cores
+        )
+
+    return tf.distribute.TPUStrategy(
+        cluster_resolver, experimental_device_assignment=device_assignment)
 
   if distribution_strategy == "multi_worker_mirrored":
     return tf.distribute.experimental.MultiWorkerMirroredStrategy(

{tf_models_nightly-2.17.0.dev20240624.dist-info → tf_models_nightly-2.17.0.dev20240625.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tf-models-nightly
-Version: 2.17.0.dev20240624
+Version: 2.17.0.dev20240625
 Summary: TensorFlow Official Models
 Home-page: https://github.com/tensorflow/models
 Author: Google Inc.

{tf_models_nightly-2.17.0.dev20240624.dist-info → tf_models_nightly-2.17.0.dev20240625.dist-info}/RECORD

@@ -1,7 +1,7 @@
 official/__init__.py,sha256=7oiypy0N82PDw9aSdcJBLVoGTd_oRSUOdvuJhMv4leQ,609
 official/common/__init__.py,sha256=Yqe5TjLrAR68lQ5G781dwlbbT6mELRByDipCtr4jQY4,610
 official/common/dataset_fn.py,sha256=MSRNhNGvMnZtjkuxFIZ90vGJgKlr2JeGaEYpz82Y3Ew,1858
-official/common/distribute_utils.py,sha256=FC4YgXohuUhp_-ZG5kNoerHdyOtGbRjCVdgKCKh5zyY,8572
+official/common/distribute_utils.py,sha256=dEciUuyBsxfv6T9rvPk_5zcvLyDFv0lPkhkFvyOGe_0,9575
 official/common/distribute_utils_test.py,sha256=ltR8wrIbfoX1miJElPQzZN2aioKX6Hx8-mfzoo-_WGc,4914
 official/common/flags.py,sha256=weXgDeMK7tZ8eY8d5LgMLRQDJYuu2hrJF-XfpOqOP34,4579
 official/common/registry_imports.py,sha256=ez-j0sOvVyYBHHGM5hNjTF0jqHUbHI0zv-4SHe_AZ4s,843
@@ -1212,9 +1212,9 @@ tensorflow_models/tensorflow_models_test.py,sha256=nc6A9K53OGqF25xN5St8EiWvdVbda
 tensorflow_models/nlp/__init__.py,sha256=4tA5Pf4qaFwT-fIFOpX7x7FHJpnyJT-5UgOeFYTyMlc,807
 tensorflow_models/uplift/__init__.py,sha256=mqfa55gweOdpKoaQyid4A_4u7xw__FcQeSIF0k_pYmI,999
 tensorflow_models/vision/__init__.py,sha256=zBorY_v5xva1uI-qxhZO3Qh-Dii-Suq6wEYh6hKHDfc,833
-tf_models_nightly-2.17.0.dev20240624.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.17.0.dev20240624.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.17.0.dev20240624.dist-info/METADATA,sha256=5-JU7SLtGBG0xywROp_9OOihPtRYoSy9THPOpqiI6Wc,1432
-tf_models_nightly-2.17.0.dev20240624.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.17.0.dev20240624.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.17.0.dev20240624.dist-info/RECORD,,
+tf_models_nightly-2.17.0.dev20240625.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.17.0.dev20240625.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.17.0.dev20240625.dist-info/METADATA,sha256=Qb7qjl8ts2IRPIuVgb36BPCqLxcU313XHjrUd3l6vc0,1432
+tf_models_nightly-2.17.0.dev20240625.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.17.0.dev20240625.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.17.0.dev20240625.dist-info/RECORD,,