I just tried this and see no improvement (both functions take same time to execute). Here are both the functions:
@tf.function(jit_compile=True)
def _forward(X, dX, W1, W2):
    Z1 = tf.matmul(X, tf.transpose(W1))
    dZ1 = tf.matmul(dX, tf.transpose(W1))

    A1 = tf.tanh(Z1)
    dA1 = tf.multiply(tf.expand_dims(1-tf.square(tf.tanh(Z1)), axis=1), dZ1)

    Z2 = tf.matmul(A1, tf.transpose(W2))
    dZ2 = tf.matmul(dA1, tf.transpose(W2))

    return Z2, tf.squeeze(dZ2, axis=-1)

@tf.function(jit_compile=True)
def forward(X, W1, W2):
    with tf.GradientTape(persistent=True) as tape:
        Z1 = tf.matmul(X, tf.transpose(W1))
        
        A1 = tf.tanh(Z1)
        
        Z2 = tf.matmul(A1, tf.transpose(W2))


    dZ2 = tape.gradient(Z2, X)
    return Z2, dZ2

I’ve always used pytorch for my work so I might not be doing things properly.
Thanks