Fix import order around from google.protobuf import text_format.

tensorflow_gnn/data/unigraph_test.py

@@ -21,13 +21,12 @@
 import apache_beam as beam
 from apache_beam.testing import test_pipeline
 from apache_beam.testing import util
+from google.protobuf import text_format
 import tensorflow as tf
 import tensorflow_gnn as tfgnn
 from tensorflow_gnn.data import unigraph
 from tensorflow_gnn.utils import test_utils
 
-from google.protobuf import text_format
-
 
 class TestUnigraph(tf.test.TestCase):
 

tensorflow_gnn/experimental/sampler/beam/accessors_test.py

@@ -18,15 +18,14 @@
 import apache_beam as beam
 from apache_beam.testing import util
 
+from google.protobuf import text_format
 import numpy as np
 import tensorflow as tf
 
 from tensorflow_gnn.experimental import sampler
 from tensorflow_gnn.experimental.sampler.beam import accessors  # pylint: disable=unused-import
 from tensorflow_gnn.experimental.sampler.beam import executor_lib
 
-from google.protobuf import text_format
-
 PCollection = beam.PCollection
 
 

tensorflow_gnn/experimental/sampler/beam/edge_samplers_test.py

@@ -22,6 +22,7 @@
 import apache_beam as beam
 from apache_beam.testing import util
 
+from google.protobuf import text_format
 import numpy as np
 import tensorflow as tf
 
@@ -30,8 +31,6 @@
 from tensorflow_gnn.experimental.sampler.beam import edge_samplers  # pylint: disable=unused-import
 from tensorflow_gnn.experimental.sampler.beam import executor_lib
 
-from google.protobuf import text_format
-
 PCollection = beam.PCollection
 
 

tensorflow_gnn/experimental/sampler/beam/executor_lib_test.py

@@ -22,14 +22,13 @@
 from apache_beam.testing import util
 from apache_beam.typehints import trivial_inference
 
+from google.protobuf import text_format
 import numpy as np
 import tensorflow as tf
 
 from tensorflow_gnn.experimental import sampler
 from tensorflow_gnn.experimental.sampler.beam import executor_lib
 
-from google.protobuf import text_format
-
 PCollection = beam.PCollection
 
 

tensorflow_gnn/experimental/sampler/beam/sampler.py

@@ -29,6 +29,7 @@
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.options.pipeline_options import SetupOptions
 from apache_beam.runners.portability import fn_api_runner
+from google.protobuf import text_format
 import tensorflow as tf
 import tensorflow_gnn as tfgnn
 from tensorflow_gnn.data import unigraph
@@ -40,7 +41,6 @@
 from tensorflow_gnn.proto import graph_schema_pb2
 import tensorflow_gnn.sampler as sampler_lib
 
-from google.protobuf import text_format
 # Placeholder for Google-internal Beam runner
 
 

tensorflow_gnn/experimental/sampler/beam/sampler_test.py

@@ -17,13 +17,12 @@
 from __future__ import annotations
 
 from absl.testing import parameterized
+from google.protobuf import text_format
 import tensorflow as tf
 import tensorflow_gnn
 from tensorflow_gnn.experimental.sampler.beam import sampler
 import tensorflow_gnn.sampler as sampler_lib
 
-from google.protobuf import text_format
-
 
 class TestGetSamplingModel(tf.test.TestCase, parameterized.TestCase):
 

tensorflow_gnn/experimental/sampler/beam/unigraph_utils_test.py

@@ -17,15 +17,14 @@
 import apache_beam as beam
 from apache_beam.testing import test_pipeline
 from apache_beam.testing import util
+from google.protobuf import text_format
 import numpy as np
 import tensorflow as tf
 import tensorflow_gnn as tfgnn
 from tensorflow_gnn.data import unigraph
 from tensorflow_gnn.experimental.sampler.beam import unigraph_utils
 from tensorflow_gnn.utils import test_utils
 
-from google.protobuf import text_format
-
 
 _CUSTOMER_IDS = b"""
   1876448 1372437 1368305 1974494 1257724 1758057 1531660 1489311 1407706

tensorflow_gnn/experimental/sampler/beam/utils_test.py

@@ -20,13 +20,12 @@
 from absl.testing import parameterized
 import apache_beam as beam
 from apache_beam.testing import util
+from google.protobuf import text_format
 import numpy as np
 import tensorflow as tf
 from tensorflow_gnn.experimental.sampler import proto as pb
 from tensorflow_gnn.experimental.sampler.beam import utils
 
-from google.protobuf import text_format
-
 PCollection = beam.PCollection
 
 rt = tf.ragged.constant

tensorflow_gnn/graph/schema_utils.py

@@ -16,6 +16,7 @@
 
 from typing import Any, Iterator, Mapping, Optional, Text, Tuple, Union
 
+from google.protobuf import text_format
 import tensorflow as tf
 from tensorflow_gnn.graph import adjacency
 from tensorflow_gnn.graph import graph_constants as gc
@@ -24,8 +25,6 @@
 from tensorflow_gnn.keras import keras_tensors as kt
 import tensorflow_gnn.proto.graph_schema_pb2 as schema_pb2
 
-from google.protobuf import text_format
-
 
 def parse_schema(schema_text: Union[bytes, str]) -> schema_pb2.GraphSchema:
   """Parse a schema from text-formatted protos.

tensorflow_gnn/graph/schema_validation_test.py

@@ -18,6 +18,7 @@
 from typing import List
 
 from absl import logging
+from google.protobuf import text_format
 import mock
 import tensorflow as tf
 
@@ -27,8 +28,6 @@
 from tensorflow_gnn.graph import schema_validation as sv
 import tensorflow_gnn.proto.graph_schema_pb2 as schema_pb2
 
-from google.protobuf import text_format
-
 
 as_tensor = tf.convert_to_tensor
 

tensorflow_gnn/sampler/sampling_spec_builder_test.py

@@ -17,12 +17,11 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 
+from google.protobuf import text_format
 import tensorflow_gnn.proto.graph_schema_pb2 as schema_pb2
 from tensorflow_gnn.sampler import sampling_spec_builder
 from tensorflow_gnn.sampler import sampling_spec_pb2
 
-from google.protobuf import text_format
-
 
 def get_schema(edge_sets=('AA', 'AB', 'AC', 'BC', 'CD'),
                add_readout_artifacts=False) -> schema_pb2.GraphSchema:

tensorflow_gnn/utils/test_utils.py

@@ -19,6 +19,7 @@
 from typing import Any
 
 from google.protobuf import text_format
+
 # Google-internal import(s).