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)