From 8809428af61ce749af27efbb4bfa9371315a46c6 Mon Sep 17 00:00:00 2001
From: Yating Jing <ytjing@google.com>
Date: Fri, 15 Sep 2023 15:59:32 +0000
Subject: [PATCH] set boolean hparam domains to be [True, False] and improve
 test

---
 .../plugins/hparams/backend_context.py        |  2 +-
 .../plugins/hparams/backend_context_test.py   | 49 +++++++++++++------
 2 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/tensorboard/plugins/hparams/backend_context.py b/tensorboard/plugins/hparams/backend_context.py
index 5c5d599beb..1267c6947e 100644
--- a/tensorboard/plugins/hparams/backend_context.py
+++ b/tensorboard/plugins/hparams/backend_context.py
@@ -335,8 +335,8 @@ def _compute_hparam_info_from_values(self, name, values):
             result.differs = len(distinct_string_values) > 1
 
         if result.type == api_pb2.DATA_TYPE_BOOL:
+            result.domain_discrete.extend([True, False])
             distinct_bool_values = set(v.bool_value for v in values)
-            result.domain_discrete.extend(distinct_bool_values)
             result.differs = len(distinct_bool_values) > 1
 
         if result.type == api_pb2.DATA_TYPE_FLOAT64:
diff --git a/tensorboard/plugins/hparams/backend_context_test.py b/tensorboard/plugins/hparams/backend_context_test.py
index 7e3b175490..67c8db4257 100644
--- a/tensorboard/plugins/hparams/backend_context_test.py
+++ b/tensorboard/plugins/hparams/backend_context_test.py
@@ -224,26 +224,31 @@ def test_experiment_with_experiment_tag_include_metrics(self):
     def test_experiment_with_session_tags(self):
         self.session_1_start_info_ = """
             hparams: [
-              {key: 'batch_size' value: {number_value: 100}},
+              {key: 'batch_size' value: {number_value: 1024}},
               {key: 'eval.timeout' value: {bool_value: false}},
               {key: 'lr' value: {number_value: 0.01}},
-              {key: 'model_type' value: {string_value: 'CNN'}}
+              {key: 'model_type' value: {string_value: 'CNN'}},
+              {key: 'optimizer_type' value: {string_value: 'momentum'}},
+              {key: 'use_batch_norm' value: {bool_value: true}}
             ]
         """
         self.session_2_start_info_ = """
             hparams:[
-              {key: 'batch_size' value: {number_value: 200}},
+              {key: 'batch_size' value: {number_value: 1024}},
               {key: 'eval.timeout' value: {bool_value: false}},
               {key: 'lr' value: {number_value: 0.02}},
-              {key: 'model_type' value: {string_value: 'LATTICE'}}
+              {key: 'model_type' value: {string_value: 'LATTICE'}},
+              {key: 'optimizer_type' value: {string_value: 'momentum'}}
             ]
         """
         self.session_3_start_info_ = """
             hparams:[
-              {key: 'batch_size' value: {number_value: 300}},
+              {key: 'batch_size' value: {number_value: 1024}},
               {key: 'eval.timeout' value: {bool_value: false}},
               {key: 'lr' value: {number_value: 0.05}},
-              {key: 'model_type' value: {string_value: 'CNN'}}
+              {key: 'model_type' value: {string_value: 'CNN'}},
+              {key: 'optimizer_type' value: {string_value: 'momentum'}},
+              {key: 'use_batch_norm' value: {bool_value: false}}
             ]
         """
         expected_exp = """
@@ -251,16 +256,16 @@ def test_experiment_with_session_tags(self):
               name: 'batch_size'
               type: DATA_TYPE_FLOAT64
               domain_interval {
-                min_value: 100.0
-                max_value: 300.0
+                min_value: 1024
+                max_value: 1024
               }
-              differs: true
+              differs: false
             }
             hparam_infos: {
               name: 'eval.timeout'
               type: DATA_TYPE_BOOL
               domain_discrete: {
-                values: [{bool_value: false}]
+                values: [{bool_value: true}, {bool_value: false}]
               }
               differs: false
             }
@@ -282,6 +287,22 @@ def test_experiment_with_session_tags(self):
               }
               differs: true
             }
+            hparam_infos: {
+              name: 'optimizer_type'
+              type: DATA_TYPE_STRING
+              domain_discrete: {
+                values: [{string_value: 'momentum'}]
+              }
+              differs: false
+            }
+            hparam_infos: {
+              name: 'use_batch_norm'
+              type: DATA_TYPE_BOOL
+              domain_discrete: {
+                values: [{bool_value: true}, {bool_value: false}]
+              }
+              differs: true
+            }
             metric_infos: {
               name: {group: '', tag: 'accuracy'}
             }
@@ -366,12 +387,12 @@ def test_experiment_with_session_tags_different_hparam_types(self):
     def test_experiment_with_session_tags_bool_types(self):
         self.session_1_start_info_ = """
             hparams:[
-              {key: 'batch_size' value: {bool_value: true}}
+              {key: 'use_batch_norm' value: {bool_value: true}}
             ]
         """
         self.session_2_start_info_ = """
             hparams:[
-              {key: 'batch_size' value: {bool_value: true}}
+              {key: 'use_batch_norm' value: {bool_value: true}}
             ]
         """
         self.session_3_start_info_ = """
@@ -380,10 +401,10 @@ def test_experiment_with_session_tags_bool_types(self):
         """
         expected_exp = """
             hparam_infos: {
-              name: 'batch_size'
+              name: 'use_batch_norm'
               type: DATA_TYPE_BOOL
               domain_discrete: {
-                values: [{bool_value: true}]
+                values: [{bool_value: true}, {bool_value: false}]
               }
               differs: false
             }