Add instruction pointer entropy metric.

Work in-progress.

Author: dan-zheng

core/lib/metrics.py

@@ -161,6 +161,33 @@ def instruction_pointers_to_images(instruction_pointer, multidevice: bool):
   return jnp.array(instruction_pointer_image_list)
 
 
+def instruction_pointers_to_entropy(instruction_pointer, multidevice: bool):
+  """Converts the given batched instruction pointer to an entropy value.
+
+  The entropy value measures the sharpness of the instruction pointer, i.e. how
+  hard vs soft it is.
+  """
+  if multidevice:
+    # instruction_pointer: device, batch_size / device, timesteps, num_nodes
+    instruction_pointer = instruction_pointer[0]
+
+  # instruction_pointer: batch_size / device, timesteps, num_nodes
+  instruction_pointer = jnp.transpose(instruction_pointer[:, :16, :],
+                                      (1, 2, 0))
+  # instruction_pointer: logging_slice_size, num_nodes, timesteps
+  instruction_pointer_image_list = [
+      instruction_pointer_to_image(ip)
+      for ip in instruction_pointer
+  ]
+  instruction_pointer_image_leading_dim_max = max(
+      image.shape[0] for image in instruction_pointer_image_list)
+  instruction_pointer_image_list = [
+      pad(image, instruction_pointer_image_leading_dim_max)
+      for image in instruction_pointer_image_list
+  ]
+  return jnp.array(instruction_pointer_image_list)
+
+
 def pad(array, leading_dim_size: int):
   """Pad the leading dimension of the given array."""
   leading_dim_difference = max(0, leading_dim_size - array.shape[0])

core/lib/trainer.py

@@ -378,6 +378,14 @@ def run_train(self, dataset_path=DEFAULT_DATASET_PATH, split='train', steps=None
             transform_fn=functools.partial(
                 metrics.instruction_pointers_to_images,
                 multidevice=config.multidevice))
+        metrics.write_metric(
+            EvaluationMetric.INSTRUCTION_POINTER_ENTROPY.value,
+            aux,
+            train_writer.scalar,
+            step,
+            transform_fn=functools.partial(
+                metrics.instruction_pointers_to_entropy,
+                multidevice=config.multidevice))
 
         # Write validation metrics.
         valid_writer.scalar('loss', valid_loss, step)