0

I have an array output and a id subject_ids.

output = [[[1, 2, 3]], [[4, 5, 6]], [[7, 8, 9]]]

subject_ids = [[0, 1], [1, 2], [0, 2]]

The numbers in ID represent the start and end positions respectively, and then I want to get the vector between them according to the start and end positions.

For example,I should get [[1, 2, 3], [4, 5, 6]] and [[4, 5, 6], [7, 8, 9]] and [[1, 2, 3], [4, 5, 6], [7, 8, 9]] in this case.

What should I do? I tried tf.slice and tf.gather, but it didn't seem to work.

AloneTogether
  • 25,814
  • 5
  • 20
  • 39
wwwxinyu
  • 11
  • 3

2 Answers2

1

If you want to use Tensorflow only, try combining tf.gather with tf.range and tf.ragged.stack:

import tensorflow as tf

output = tf.constant([
                      [[1, 2, 3]], 
                      [[4, 5, 6]], 
                      [[7, 8, 9]]
                      ])

subject_ids = tf.constant([[0, 1], [1, 2], [0, 2]])

ragged_ouput = tf.ragged.stack([tf.gather(output, tf.range(subject_ids[i, 0], subject_ids[i, 1] + 1)) for i in tf.range(0, tf.shape(subject_ids)[0])], axis=0)
ragged_ouput = tf.squeeze(ragged_ouput, axis=2)
print(ragged_ouput)

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

Update 1:

import tensorflow as tf
tf.config.run_functions_eagerly(True)

output = tf.constant([
                      [[1, 2, 3]], 
                      [[4, 5, 6]], 
                      [[7, 8, 9]]
                      ])

subject_ids = tf.constant([[0, 1], [1, 2], [0, 2]])

def slice_tensor(x):
  return tf.ragged.stack([tf.gather(output, tf.range(x[0], x[1] + 1))], axis=0)

ragged_ouput = tf.map_fn(slice_tensor, subject_ids, fn_output_signature=tf.RaggedTensorSpec(shape=[1, None, None, 3],
                                                                    dtype=tf.int32,
                                                                    ragged_rank=2,
                                                                    row_splits_dtype=tf.int64))
ragged_ouput = tf.squeeze(ragged_ouput, axis=1)
tf.print(ragged_ouput, summarize=-1)

[[[[1, 2, 3]], [[4, 5, 6]]], [[[4, 5, 6]], [[7, 8, 9]]], [[[1, 2, 3]], [[4, 5, 6]], [[7, 8, 9]]]]
AloneTogether
  • 25,814
  • 5
  • 20
  • 39
  • sorry, I run this but got the `TypeError: Tensor objects are only iterable when eager execution is enabled. To iterate over this tensor use tf.map_fn.` Can't loops be used in tensorflow's model? – wwwxinyu Jan 03 '22 at 02:45
  • 1
    Of course, they can be used, but it depends where. The problem is you show two lines of code in your question without mentioning anything else and you expect a complete working example which will fit into the rest of your existing code (which none of us can see). Obviously, it might not work, because you have failed to provide further details. Anyway, I updated my answer with an example with tf.map_fn – AloneTogether Jan 03 '22 at 10:16
  • Yes, I wanted to simplify the problem at first, so I didn't post all the code. Thanks very much ! That helps me a lot. – wwwxinyu Jan 03 '22 at 11:00
0

How about just

>>> [output[np.arange(x, y+1)] for x, y in subject_ids] 
[array([[[1, 2, 3]],
        [[4, 5, 6]]]),
        
 array([[[4, 5, 6]],
        [[7, 8, 9]]]),
        
 array([[[1, 2, 3]],
        [[4, 5, 6]],
        [[7, 8, 9]]])]
  • sorry, I run this but also got the `TypeError: Tensor objects are only iterable when eager execution is enabled. To iterate over this tensor use tf.map_fn.` Can't loops be used in tensorflow's model? – wwwxinyu Jan 03 '22 at 02:48