def batch(self, dataset, batch_size):
def _batch(dataset, batch_size):
return dataset.batch(batch_size)
batch_fn_map = {"graph": _batch}
return diverse_batch(
dataset, batch_size, batch_fn_map,
default_batch_fn=self.node_embedder.batch)
def batch(self, dataset, batch_size):
batch_fn_map = {"features": shifted_batch}
return diverse_batch(
dataset,
batch_size,
batch_fn_map,
default_batch_fn=lambda dataset, batch_size: dataset.padded_batch(
batch_size, get_padded_shapes(dataset)))
def batch(self, dataset, batch_size):
"""
"""
def _batch(dataset, batch_size):
return dataset.padded_batch(batch_size, get_padded_shapes(dataset))
batch_fn_map = {("primary_path", "primary_path_length"): _batch}
return diverse_batch(dataset,
batch_size,
batch_fn_map,
default_batch_fn=self.graph_inputter.batch)
def batch_fn(dataset, batch_size):
return diverse_batch(dataset, batch_size,
(feat_batch_fn, labels_batch_fn))
def batch(self, dataset, batch_size):
batch_fn_map = {"features": shifted_batch}
return diverse_batch(dataset, batch_size, batch_fn_map)
def batch(self, dataset, batch_size):
def _padded_batch(dataset, batch_size):
return dataset.padded_batch(batch_size, get_padded_shapes(dataset))
batch_fn_map = {("features", "length"): _padded_batch}
return diverse_batch(dataset, batch_size, batch_fn_map)