fix(integrations): make Keras WandbCallback compatible with TF version >= 2.11.0

[ayulockin]

Fixes WB-11561

Description

In TensorFlow 2.9, an experimental version of the new Keras Optimizer API, tf.keras.optimizers.experimental was released. It was done to provide a more unified and expanded catalog of built-in optimizers which can be more easily customized and extended.

In TensorFlow 2.11, these new optimizers are on by default.

Our Keras WandbCallback, log_gradient feature creates a _CustomOptimizer by subclassing old optimizers API.

This PR creates the _CustomOptimizer using the correct parent class. This solves the problem of log_gradient feature not working with TF 2.11.0.

Testing

This was manually tested with TF version 2.8.0, 2.9.0, 2.10.0, and 2.11.0. The W&B run pages are linked below:

• 2.8.0
• 2.9.0
• 2.10.0
• 2.11.0

Checklist

• Include reference to internal ticket "Fixes WB-NNNN" and/or GitHub issue "Fixes #NNNN" (if applicable)
• Ensure PR title compliance with the conventional commits standards

[codecov]

Codecov Report

❗ No coverage uploaded for pull request base (main@8c7d8b3). Click here to learn what that means.
The diff coverage is 100.00%.

❗ Current head 137fa3f differs from pull request most recent head 37f5a9f. Consider uploading reports for the commit 37f5a9f to get more accurate results

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #4533   +/-   ##
=======================================
  Coverage        ?   83.03%           
=======================================
  Files           ?      261           
  Lines           ?    33227           
  Branches        ?        0           
=======================================
  Hits            ?    27590           
  Misses          ?     5637           
  Partials        ?        0           

Flag	Coverage Δ
functest	56.66% <100.00%> (?)
unittest	73.30% <100.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
wandb/integration/keras/keras.py	73.76% <100.00%> (ø)

[moredatarequired]

Thank you! Could you also roll back the changes from #4508 to ensure that 2.11 gets tested in CI?

[dmitryduev]

Excellent, thanks a lot for the fix @ayulockin!
+1 on what @moredatarequired requested and then should be good to go in.

[dmitryduev]

mypy issue seems unrelated, will look into that separately.

[dmitryduev]

@ayulockin @moredatarequired I merged a revert of #4508 into this branch.