diff --git a/tensorflow/python/keras/callbacks.py b/tensorflow/python/keras/callbacks.py
index 2deca3dd414..f3169d2081c 100644
--- a/tensorflow/python/keras/callbacks.py
+++ b/tensorflow/python/keras/callbacks.py
@@ -2201,7 +2201,7 @@ class TensorBoard(Callback, version_utils.TensorBoardVersionSelector):
         train_fn = self.model.train_function
         # If the train_function is a `tf.function`, we can write out a graph
         if hasattr(train_fn, 'function_spec'):
-          summary_ops_v2.graph(train_fn._concrete_stateful_fn.graph, step=0)  # pylint: disable=protected-access
+          summary_ops_v2.graph(train_fn._concrete_stateful_fn.graph)  # pylint: disable=protected-access
 
   def _write_keras_model_summary(self):
     """Writes Keras graph network summary to TensorBoard."""
diff --git a/tensorflow/python/keras/callbacks_v1.py b/tensorflow/python/keras/callbacks_v1.py
index 4e763e00761..79a8e9d2efc 100644
--- a/tensorflow/python/keras/callbacks_v1.py
+++ b/tensorflow/python/keras/callbacks_v1.py
@@ -170,7 +170,7 @@ class TensorBoard(callbacks.TensorBoard):
       self.writer = summary_ops_v2.create_file_writer_v2(self.log_dir)
       if not model.run_eagerly and self.write_graph:
         with self.writer.as_default():
-          summary_ops_v2.graph(K.get_graph(), step=0)
+          summary_ops_v2.graph(K.get_graph())
     elif self.write_graph:
       self.writer = tf_summary.FileWriter(self.log_dir, K.get_graph())
     else:
diff --git a/tensorflow/python/kernel_tests/summary_ops_test.py b/tensorflow/python/kernel_tests/summary_ops_test.py
index bce1e0d9718..e64e75bde0d 100644
--- a/tensorflow/python/kernel_tests/summary_ops_test.py
+++ b/tensorflow/python/kernel_tests/summary_ops_test.py
@@ -1217,7 +1217,7 @@ class SummaryOpsTest(test_util.TensorFlowTestCase):
       summary_ops.set_step(None)
 
   @test_util.run_v2_only
-  def testGraphV2_graph(self):
+  def testGraph_graph(self):
 
     @def_function.function
     def f():
@@ -1226,13 +1226,13 @@ class SummaryOpsTest(test_util.TensorFlowTestCase):
       return x**y
 
     def summary_op_fn():
-      summary_ops.graph_v2(f.get_concrete_function().graph)
+      summary_ops.graph(f.get_concrete_function().graph)
 
     event = self.exec_summary_op(summary_op_fn)
     self.assertIsNotNone(event.graph_def)
 
   @test_util.run_v2_only
-  def testGraphV2_graphDef(self):
+  def testGraph_graphDef(self):
 
     @def_function.function
     def f():
@@ -1241,15 +1241,15 @@ class SummaryOpsTest(test_util.TensorFlowTestCase):
       return x**y
 
     def summary_op_fn():
-      summary_ops.graph_v2(f.get_concrete_function().graph.as_graph_def())
+      summary_ops.graph(f.get_concrete_function().graph.as_graph_def())
 
     event = self.exec_summary_op(summary_op_fn)
     self.assertIsNotNone(event.graph_def)
 
   @test_util.run_v2_only
-  def testGraphV2_invalidData(self):
+  def testGraph_invalidData(self):
     def summary_op_fn():
-      summary_ops.graph_v2('hello')
+      summary_ops.graph('hello')
 
     with self.assertRaisesRegex(
         ValueError,
@@ -1258,7 +1258,7 @@ class SummaryOpsTest(test_util.TensorFlowTestCase):
       self.exec_summary_op(summary_op_fn)
 
   @test_util.run_v2_only
-  def testGraphV2_fromGraphMode(self):
+  def testGraph_fromGraphMode(self):
 
     @def_function.function
     def f():
@@ -1268,7 +1268,7 @@ class SummaryOpsTest(test_util.TensorFlowTestCase):
 
     @def_function.function
     def g(graph):
-      summary_ops.graph_v2(graph)
+      summary_ops.graph(graph)
 
     def summary_op_fn():
       graph_def = f.get_concrete_function().graph.as_graph_def(add_shapes=True)
diff --git a/tensorflow/python/ops/summary_ops_v2.py b/tensorflow/python/ops/summary_ops_v2.py
index 83cc09e443c..6be1ea04968 100644
--- a/tensorflow/python/ops/summary_ops_v2.py
+++ b/tensorflow/python/ops/summary_ops_v2.py
@@ -451,7 +451,7 @@ def initialize(
   if graph is not None:
     data = _serialize_graph(graph)
     x = array_ops.placeholder(dtypes.string)
-    session.run(_graph(x, 0), feed_dict={x: data})
+    session.run(graph_v1(x, 0), feed_dict={x: data})
 
 
 @tf_export("summary.create_file_writer", v1=[])
@@ -966,8 +966,7 @@ def audio(name, tensor, sample_rate, max_outputs, family=None, step=None):
   return summary_writer_function(name, tensor, function, family=family)
 
 
-# TODO(b/171925996): rename this to `graph_v1` after changing all callsites.
-def graph(param, step=None, name=None):
+def graph_v1(param, step=None, name=None):
   """Writes a TensorFlow graph to the summary interface.
 
   The graph summary is, strictly speaking, not a summary. Conditions
@@ -1012,13 +1011,8 @@ def graph(param, step=None, name=None):
         writer._resource, _choose_step(step), tensor, name=name)  # pylint: disable=protected-access
 
 
-_graph = graph  # for functions with a graph parameter
-
-
-# TODO(b/171925996): rename this to `graph` for consistency when graph_v1 is
-# renamed.
 @tf_export("summary.graph", v1=[])
-def graph_v2(graph_data):
+def graph(graph_data):
   """Writes a TensorFlow graph summary.
 
   Write an instance of `tf.Graph` or `tf.compat.v1.GraphDef` as summary only