Does shap's gradientexplainer take in a tensorflow function or a model?

I'm using tensorflow 2 trying to call the gradientexplainer. I'm looking at the docs and it states the following:

Note that for TensowFlow 2 you must pass a tensorflow function, not a tuple of input/output tensors

However it seems that you need to pass in your model directly, you can't pass a tensorflow function.

My model expects a dictionary instead of a nparray or dataframe, so I am attempting to call gradientexplainer like so:

@tf.functiondef model_fcn(data):    feed_dict = {k: tf.convert_to_tensor(data[:, i:i+1]) for i, k in enumerate(feature_keys)}    return model(feed_dict)gradientexplainer = shap.GradientExplainer(model_fcn, X)

However I get this error: ValueError: <class tensorflow.python.eager.polymorphic_function.polymorphic_function.Function'> is not currently a supported model type!

Just wondering if you can pass in a function to gradientexplainer or if it has to be the keras model