AttributeError: module 'tensorflow_addons.image' has no attribute 'gradients'

Question

I am working on a project on google colab where I need to visualize gradient maps for an image using TensorFlow and TensorFlow Addons. I have defined a function plot_gradient_maps that uses tfa.image.gradients to compute gradients. However, I am encountering an AttributeError stating that the module tensorflow_addons.image has no attribute gradients.

Here is the relevant part of the code:

import tensorflow_addons as tfa

def input_img(path): #read source image file
    image = tf.image.decode_png(tf.io.read_file(path))
    image = tf.expand_dims(image, axis=0)
    image = tf.cast(image, tf.float32)
    image = tf.image.resize(image, [160,160])
    return image
def normalize_image(img): #normalise image
    grads_norm = img[:,:,0]+ img[:,:,1]+ img[:,:,2]
    grads_norm = (grads_norm - tf.reduce_min(grads_norm))/ (tf.reduce_max(grads_norm)- tf.reduce_min(grads_norm))
    return grads_norm

def plot_maps(img1, img2,vmin=0.3,vmax=0.7, mix_val=2): 
    f = plt.figure(figsize=(15,45))
    plt.subplot(1,3,1)
    plt.imshow(img1,vmin=vmin, vmax=vmax, cmap="terrain")
    plt.title('Saliency map')
    plt.axis("off")
    plt.subplot(1,3,2)
    plt.imshow(img2, cmap = "terrain")
    plt.title('Input image')
    plt.axis("off")
    plt.subplot(1,3,3)
    plt.imshow(img1*mix_val+img2/mix_val, cmap = "terrain" )
    plt.title('Overlayed image')
    plt.axis("off")
def plot_orig_img(filepath):
    img_path = filepath
    input_im = input_img(img_path)
    input_im = tf.keras.applications.efficientnet_v2.preprocess_input(input_im)
    plt.imshow(normalize_image(input_im[0]))
    plt.axis('off');
    return input_im
def plot_gradient_maps(input_im): 
    with tf.GradientTape() as tape:
        tape.watch(input_im)
        result_img = model(input_im)
        max_idx = tf.argmax(result_img, axis=1)
        max_score = result_img[0, max_idx[0]]

    # Use tfa.image.gradients 
    grads = tfa.image.gradients(max_score, input_im) 
    plot_maps(normalize_image(grads[0]), normalize_image(input_im[0]))


input_im = plot_orig_img(filepath="/content/work_directory/test/glaucoma/1209_right.jpg");
plot_gradient_maps(input_im)

Error:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-121-26a30edde128> in <cell line: 1>()
----> 1 plot_gradient_maps(input_im)

<ipython-input-119-b5bf95270540> in plot_gradient_maps(input_im)
     48 
     49     # Use tfa.image.gradients
---> 50     grads = tfa.image.gradients(max_score, input_im)
     51     plot_maps(normalize_image(grads[0]), normalize_image(input_im[0]))
     52 

AttributeError: module 'tensorflow_addons.image' has no attribute 'gradients'

The gradients function seems to be missing from the tfa.image module. I have checked the documentation, but I couldn't find any information on this specific issue.

I attempted to use tfa.image.gradients for computing gradients, following the examples and documentation available. However, the AttributeError mentioned above was encountered.

I expected the code to execute without errors and generate gradient maps using the tfa.image.gradients function.

How to correctly use tfa.image.gradients? or suggest an alternative approach for computing gradients in TensorFlow Addons.