This fix adjusts import to avoid accidentally pulling in keras in clu. #336

PiperOrigin-RevId: 617332775

Author: CLU Authors

clu/metric_writers/tf/summary_writer.py

@@ -27,7 +27,7 @@
 from clu.metric_writers import interface
 import tensorflow as tf
 
-from tensorboard.plugins.hparams import api as hparams_api
+from tensorboard.plugins.hparams import summary_v2
 
 
 Array = interface.Array
@@ -95,7 +95,7 @@ def write_histograms(
 
   def write_hparams(self, hparams: Mapping[str, Any]):
     with self._summary_writer.as_default():
-      hparams_api.hparams(dict(utils.flatten_dict(hparams)))
+      summary_v2.hparams(dict(utils.flatten_dict(hparams)))
 
   def flush(self):
     self._summary_writer.flush()