partition or mask or filter a Tensor in tensorflow.js

I have 2 Tensors of same length, data and groupIds. I want to split data into several groups by the corresponding values in groupId. For example,

const data = tf.tensor([1,2,3,4,5]);
const groupIds = tf.tensor([0,1,1,0,0]);
// expected result: [tf.tensor([1,4,5]), tf.tensor([2,3])]

In Tensorflow there is tf.dynamic_partition which does exactly that. Tensorflow.js doesn't seem to have a similar method. I also looked into mask or filtering as work-arounds, but they don't exist either. Does anyone have an idea how to implement this?

To partition your tensor, you can first iterate over your ids tensor to get the number of subtensor to create and the index of the elements it should contain. This information can be stored in an object where the key is the number of the partition in the ids array and the value is an array of indexes.

const data = tf.tensor([6,2,8,4,5]);
const ids = tf.tensor([0,1,1,0,2]);

const data2 = tf.tensor([[6,2],[8,4], [5, 4], [6, 5]]);
const ids2 = tf.tensor([0,1,1,0]);

const filterT = (t, p) => {
  t.print()
  p.print()
  const l = p.unstack().reduce((a, b, i) => {
  const v = b.dataSync()[0]
  if (Object.keys(a).includes(v.toString())) {
    a[v].push(i)
  } else {
    a[v] = [i]
  }
  return a
}, {})

  const r = Object.keys(l).map(k => t.gather(tf.tensor1d(l[k], 'int32')))
  r.forEach(e => e.print())
}

filterT(data, ids)
filterT(data2, ids2)

<html>
  <head>
    <!-- Load TensorFlow.js -->
    <script src="https://cdnjs.cloudflare.com/ajax/libs/tensorflow/0.12.4/tf.js"> </script>
  </head>

  <body>
  </body>
</html>

来源：https://stackoverflow.com/questions/51218420/partition-or-mask-or-filter-a-tensor-in-tensorflow-js

标签

tensorflow

tensorflow.js