Hello
a suggestion 
import clr
import sys
import System
clr.AddReference('RevitAPI')
import Autodesk
from Autodesk.Revit.DB import *
import Autodesk.Revit.DB as DB
clr.AddReference('RevitServices')
import RevitServices
from RevitServices.Persistence import DocumentManager
from RevitServices.Transactions import TransactionManager
doc = DocumentManager.Instance.CurrentDBDocument
elements = UnwrapElement(IN[0]) if isinstance(IN[0],list) else [UnwrapElement(IN[0])]
dicDocLink = {link.GetLinkDocument().GetHashCode() : link for link in FilteredElementCollector(doc).OfClass(RevitLinkInstance) }
OUT = [dicDocLink.get(elem.Document.GetHashCode()) for elem in elements]