From 83360c55f5afe0959b996e8653a9e2c85a310821 Mon Sep 17 00:00:00 2001
From: OptaxDev <no-reply@google.com>
Date: Mon, 11 Dec 2023 15:01:12 -0800
Subject: [PATCH] Fix typing issue in prodigy.py for python 3.9

PiperOrigin-RevId: 589961288
---
 optax/contrib/prodigy.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/optax/contrib/prodigy.py b/optax/contrib/prodigy.py
index a6c5bab2..6b02271a 100644
--- a/optax/contrib/prodigy.py
+++ b/optax/contrib/prodigy.py
@@ -46,7 +46,7 @@ class ProdigyState(NamedTuple):
 def prodigy(
     learning_rate: base.ScalarOrSchedule = 0.1,
     betas: tuple[float, float] = (0.9, 0.999),
-    beta3: float | None = None,
+    beta3: Optional[float] = None,
     eps: float = 1e-8,
     estim_lr0: float = 1e-6,
     estim_lr_coef: float = 1.0,