Tensorflow - reshape or split tensor on pre-defined value

Question

I need to define a graph operation that splits a 2D tensorflow tensor into a 3D one based on a specific value (denoted as X here), and padded on the second and third dimension.

For example, say we have the following tensor:

[[1, 2, 3, X, 4, 5, 6],
 [1, 2, X, 5, 6, X, 8],
 [1, 2, 5, X, 8, 9, P]]

• X being the to split on (it could be equal to 0 or anything else).
• P is the symbol for padding.

The goal is to have the following:

[[[1, 2, 3], [4, 5, 6], [P, P, P]],
[[[1, 2, P], [5, 6, P], [8, P, P]],
[[[1, 2, 5], [8, 9, P], [P, P, P]]

Tensorflow provides a split function, but it does not solve this problem. The reshape function is not the solution either as the split needs to based on a specific value.

Thank you.

Answer 1

Maybe you can use tf.py_func to do the split operation.