Update jax.host_id() and jax.host_count() to jax.process_index() and jax.process_count(), respectively.

PiperOrigin-RevId: 863032172

Author: danielsuo

vit_jax_augreg.ipynb

@@ -4355,10 +4355,10 @@
             "INFO:absl:Selected fillename=\"R_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0\" for \"R_Ti_16\" with final_val=0.390\n",
             "INFO:absl:Found no checkpoint file at ./workdirs/1627539136\n",
             "INFO:absl:Will start/continue training at initial_step=1\n",
-            "/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:387: UserWarning: jax.host_count has been renamed to jax.process_count. This alias will eventually be removed; please update your code.\n",
-            "  \"jax.host_count has been renamed to jax.process_count. This alias \"\n",
-            "/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:374: UserWarning: jax.host_id has been renamed to jax.process_index. This alias will eventually be removed; please update your code.\n",
-            "  \"jax.host_id has been renamed to jax.process_index. This alias \"\n",
+            "/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:387: UserWarning: jax.process_count has been renamed to jax.process_count. This alias will eventually be removed; please update your code.\n",
+            "  \"jax.process_count has been renamed to jax.process_count. This alias \"\n",
+            "/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:374: UserWarning: jax.process_index has been renamed to jax.process_index. This alias will eventually be removed; please update your code.\n",
+            "  \"jax.process_index has been renamed to jax.process_index. This alias \"\n",
             "INFO:absl:Hyperparameters: {'accum_steps': 0, 'base_lr': 0.01, 'batch': 512, 'batch_eval': 120, 'checkpoint_every': 1000, 'dataset': 'tf_flowers', 'decay_type': 'cosine', 'eval_every': 100, 'grad_norm_clip': 1.0, 'model': {'classifier': 'token', 'hidden_size': 192, 'name': 'R+ViT-Ti_16', 'patches': {'size': (8, 8)}, 'representation_size': None, 'resnet': {'num_layers': (), 'width_factor': 1}, 'transformer': {'attention_dropout_rate': 0.0, 'dropout_rate': 0.0, 'mlp_dim': 768, 'num_heads': 3, 'num_layers': 12}}, 'model_or_filename': 'R_Ti_16', 'optim_dtype': 'bfloat16', 'pp': {'crop': 224, 'resize': 448, 'test': 'train[90%:]', 'train': 'train[:90%]'}, 'prefetch': 2, 'pretrained_dir': 'gs://vit_models/augreg', 'progress_every': 10, 'shuffle_buffer': 1000, 'tfds_data_dir': None, 'tfds_manual_dir': None, 'total_steps': 100, 'trainer': 'train', 'warmup_steps': 10}\n",
             "INFO:absl:Starting training loop; initial compile can take a while...\n",
             "/usr/local/lib/python3.7/dist-packages/jax/_src/profiler.py:167: UserWarning: StepTraceContext has been renamed to StepTraceAnnotation. This alias will eventually be removed; please update your code.\n",
@@ -4597,10 +4597,10 @@
             "INFO:absl:Selected fillename=\"R_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0\" for \"R_Ti_16\" with final_val=0.390\n",
             "INFO:absl:Found no checkpoint file at ./workdirs/1627539429\n",
             "INFO:absl:Will start/continue training at initial_step=1\n",
-            "/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:387: UserWarning: jax.host_count has been renamed to jax.process_count. This alias will eventually be removed; please update your code.\n",
-            "  \"jax.host_count has been renamed to jax.process_count. This alias \"\n",
-            "/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:374: UserWarning: jax.host_id has been renamed to jax.process_index. This alias will eventually be removed; please update your code.\n",
-            "  \"jax.host_id has been renamed to jax.process_index. This alias \"\n",
+            "/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:387: UserWarning: jax.process_count has been renamed to jax.process_count. This alias will eventually be removed; please update your code.\n",
+            "  \"jax.process_count has been renamed to jax.process_count. This alias \"\n",
+            "/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:374: UserWarning: jax.process_index has been renamed to jax.process_index. This alias will eventually be removed; please update your code.\n",
+            "  \"jax.process_index has been renamed to jax.process_index. This alias \"\n",
             "INFO:absl:Hyperparameters: {'accum_steps': 0, 'base_lr': 0.01, 'batch': 512, 'batch_eval': 120, 'checkpoint_every': 1000, 'dataset': './flower_photos', 'decay_type': 'cosine', 'eval_every': 100, 'grad_norm_clip': 1.0, 'model': {'classifier': 'token', 'hidden_size': 192, 'name': 'R+ViT-Ti_16', 'patches': {'size': (8, 8)}, 'representation_size': None, 'resnet': {'num_layers': (), 'width_factor': 1}, 'transformer': {'attention_dropout_rate': 0.0, 'dropout_rate': 0.0, 'mlp_dim': 768, 'num_heads': 3, 'num_layers': 12}}, 'model_or_filename': 'R_Ti_16', 'optim_dtype': 'bfloat16', 'pp': {'crop': 224, 'resize': 448, 'test': 'train[90%:]', 'train': 'train[:90%]'}, 'prefetch': 2, 'pretrained_dir': 'gs://vit_models/augreg', 'progress_every': 10, 'shuffle_buffer': 1000, 'tfds_data_dir': None, 'tfds_manual_dir': None, 'total_steps': 100, 'trainer': 'train', 'warmup_steps': 10}\n",
             "INFO:absl:Starting training loop; initial compile can take a while...\n",
             "/usr/local/lib/python3.7/dist-packages/jax/_src/profiler.py:167: UserWarning: StepTraceContext has been renamed to StepTraceAnnotation. This alias will eventually be removed; please update your code.\n",