Vision Transformer attention map by keypoint location - TensorFlow

I have trained a ViT model on TensorFlow for keypoint estimation based on https://github.com/yangsenius/TransPose and I would like to simulate the attention maps of each keypoint like this: https://raw.githubusercontent.com/yangsenius/TransPose/main/attention_map_image_dependency_transposeh_thres_0.00075.jpg

I have found the code on Pytorch but I have no idea about how to simulate it on TensorFlow: https://github.com/yangsenius/TransPose/blob/dab9007b6f61c9c8dce04d61669a04922bbcd148/visualize.py#L128

I have solved it by getting the output of the previous layer of the multihead attention layer and passing it by the multihead attention:

atten_maps_hooks = [Model(inputs = model.input, outputs = model.layers[getLayerIndexByName(model, 'encoded_0') - 1].output),
                        Model(inputs = model.input, outputs = model.layers[getLayerIndexByName(model, 'encoded_1') - 1].output),
                        Model(inputs = model.input, outputs = model.layers[getLayerIndexByName(model, 'encoded_2') - 1].output),
                        Model(inputs = model.input, outputs = model.layers[getLayerIndexByName(model, 'encoded_3') - 1].output),
                        Model(inputs = model.input, outputs = model.layers[getLayerIndexByName(model, 'encoded_4') - 1].output),
                        Model(inputs = model.input, outputs = model.layers[getLayerIndexByName(model, 'encoded_5') - 1].output)]

for i in range(len(atten_maps_hooks)):
      temp = atten_maps_hooks[i].predict(input)
      mha, scores = model.get_layer('encoded_' + str(i))(temp, temp, return_attention_scores = True)
      enc_atten_maps_hwhw.append(scores.numpy()[0].reshape(shape + shape))