Can tf.gradienttape() calculate gradient of other library's function

Question

If I include inside the tf.GradientTape() some functions from other Python libraries, like `sklearn.decomposition.PCA.inverse_transform()', can TensorFlow calculate gradients from that function?

Specifically, can tf automatically differetiate pca_inverse_tranform = pca.inverse_transform(h2)?

...
from sklearn.decomposition import PCA
pca = PCA(n_components=10)
pca.fit(x)
...
with tf.GradientTape() as tape:
    h1 = x@w1 + tf.broadcast_to(b1, [x.shape[0], 256])
    h1 = tf.nn.relu(h1)
    h2 = h1@w2 + tf.broadcast_to(b2, [x.shape[0], 10])
    h2 = tf.nn.relu(h2)

    pca_inverse_tranform = pca.inverse_transform(h2)

    loss = tf.square(pca_inverse_tranform - target)
    loss = tf.reduce_mean(loss)
[dl_dw1, dl_db1, dl_dw2, dl_db2] = tape.gradient(loss, [w1,b1,w2,b2])

Answer 1

I found the answer from tf docs. It says The tape can't record the gradient path if the calculation exits TensorFlow. For example:

x = tf.Variable([[1.0, 2.0],
                 [3.0, 4.0]], dtype=tf.float32)

with tf.GradientTape() as tape:
  x2 = x**2

  # This step is calculated with NumPy
  y = np.mean(x2, axis=0)

  # Like most ops, reduce_mean will cast the NumPy array to a constant tensor
  # using `tf.convert_to_tensor`.
  y = tf.reduce_mean(y, axis=0)

print(tape.gradient(y, x))

So, the answer to the question is "No, tf can't calculate gradient of other library's functions".