Fix various type annotations.#306

copybara-service[bot] wants to merge 1 commit intomaingoogle/learned_optimization:mainfrom

test_572595145google/learned_optimization:test_572595145Copy head branch name to clipboard