Equalization layer gets retraced on every call (even without any operations inside augment_image)

[quantumalaviya]

I discovered this problem with equalization while trying to add changes for my AugMix PR. It was already known that the Equalization layer gets retraced on every call. While trying to find the reason, I realized how this problem persists even when all operations inside augment_image are removed (i.e. augment_image simply returns the input). This seems to be an issue with BaseImageAugmentationLayer. The strange part is that some of the other layers I tested (FourierMix, Posterization, and Solarization) don't share this issue.

Steps to reproduce:
Edit the Equalization.augment_image function to:

    def augment_image(self, image, transformation=None):
        return image

and call the layer multiple times (I just used a loop to repeatedly called it on the same image batch)

[bhack]

Can we close this now?

[LukeWood]

No, it still gets retraced. I think we fixed the dynamic shape issue, not retracing issue (the PR was improperly titled)

We need to verify the retracing issue is fixed.

[sebastian-sz]

I think this could be related to: #241

[bhack]

@sebastian-sz Yes as we had a quite detailed analisys there I don't know if we Need also this one.

[quantumalaviya]

I think this can be closed @LukeWood

#241