please use `optimizer._learning_rate`.
    *   **You implemented a custom optimizer based on the old optimizer.**
        Please set your optimizer to subclass
        `tf.keras.optimizer.legacy.XXX`. If you want to migrate to the new
        optimizer and find it does not support your optimizer, please file
        an issue in the Keras GitHub repo.

class AbstractFunction : public core::RefCounted {
 protected:
  enum AbstractFunctionKind { kGraph, kMlir };
  explicit AbstractFunction(AbstractFunctionKind kind) : kind_(kind) {}

 public:
  // Returns which subclass is this instance of.
  AbstractFunctionKind getKind() const { return kind_; }

  // Returns the AbstractFunction as a FunctionDef.
  virtual Status GetFunctionDef(const FunctionDef**) = 0;