def _to_proto(self): pb = dataset_options_pb2.DistributeOptions() pb.auto_shard_policy = AutoShardPolicy._to_proto( self.auto_shard_policy) # pylint: disable=protected-access if self.num_devices is not None: pb.num_devices = self.num_devices return pb
def testProtoOptionsDefaultValuesRoundTrip(self): pb = dataset_options_pb2.Options() options = dataset_ops.Options() options._from_proto(pb) result = options._to_proto() expected_pb = dataset_options_pb2.Options() expected_pb.distribute_options.CopyFrom( dataset_options_pb2.DistributeOptions()) expected_pb.optimization_options.CopyFrom( dataset_options_pb2.OptimizationOptions()) expected_pb.threading_options.CopyFrom( dataset_options_pb2.ThreadingOptions()) self.assertProtoEquals(expected_pb, result)