I have run your code exp/Optimizer_2.ipynb as it is in Colab Notebook but found the following error:
ValueError: Argument learning_rate should be float, or an instance of LearningRateSchedule, or a callable (that takes in the current iteration value and returns the corresponding learning rate value). Received instead: learning_rate=MySGDOptimizer