parent.getDefiningOp<TF::TensorListSetItemOp>()) {
        parent = set_item.getInputHandle();
      } else {
        break;
      }
    }
    if (auto block_arg = parent.dyn_cast<mlir::BlockArgument>()) {
      return block_arg.getArgNumber();
    }
    // Returns -1 if we don't find which this result maps to.
    return -1;
  };

  llvm::SmallVector<Value, 4> returns;