treeflow.tf_util.vectorize module

treeflow.tf_util.vectorize.broadcast_structure(elems, event_shape, batch_shape)

treeflow.tf_util.vectorize.reshape_structure(elems, event_shape, new_batch_shape)

treeflow.tf_util.vectorize.vectorize_over_batch_dims(fn, elems, event_shape, batch_shape, vectorized_map=True, fn_output_signature=None)