| 789 | } |
| 790 | |
| 791 | async next(): Promise<IteratorResult<O>> { |
| 792 | const item = await this.upstream.next(); |
| 793 | if (item.done) { |
| 794 | return {value: null, done: true}; |
| 795 | } |
| 796 | const inputTensors = tf.tensor_util.getTensorsInContainer(item.value as {}); |
| 797 | // Careful: the transform may mutate the item in place. |
| 798 | // That's why we have to remember the input Tensors above, and then |
| 799 | // below dispose only those that were not passed through to the output. |
| 800 | // Note too that the transform function is responsible for tidying |
| 801 | // any intermediate Tensors. Here we are concerned only about the |
| 802 | // inputs. |
| 803 | const mapped = await this.transform(item.value); |
| 804 | const outputTensors = tf.tensor_util.getTensorsInContainer(mapped as {}); |
| 805 | |
| 806 | // TODO(soergel) faster intersection |
| 807 | // TODO(soergel) move to tf.disposeExcept(in, out)? |
| 808 | for (const t of inputTensors) { |
| 809 | if (!tf.tensor_util.isTensorInList(t, outputTensors)) { |
| 810 | t.dispose(); |
| 811 | } |
| 812 | } |
| 813 | return {value: mapped, done: false}; |
| 814 | } |
| 815 | } |
| 816 | |
| 817 | // Iterators that maintain a queue of pending items |