merging greptile changes again

Author: tdophung

docs/examples/quickstart_jax_utils.py

@@ -62,8 +62,8 @@ def create_train_step_fn(
     if forward_kwargs is None:
         forward_kwargs = {}
 
-    def loss_fn(variables : Any, inp: jnp.ndarray, grad_target: jnp.ndarray, dropout_key):
-        rngs = {'dropout': dropout_key}
+    def loss_fn(variables: Any, inp: jnp.ndarray, grad_target: jnp.ndarray, dropout_key):
+        rngs = {"dropout": dropout_key}
         with te.fp8_autocast(**fp8_autocast_kwargs):
             # Forward Pass: Apply the model using current parameters and variables
             call_kwargs = {**forward_kwargs, "rngs": rngs}
@@ -97,21 +97,20 @@ def create_train_step_fn_vjp(
 
     def train_step_fn(variables: Any, inp: jnp.ndarray, grad_target: jnp.ndarray, dropout_key):
         """Compute forward pass and VJP in one step"""
-        
+
         # Define forward function that closes over grad_target and dropout_key
         def forward_fn(variables: Any, inp: jnp.ndarray):
             """Pure forward function for VJP computation"""
-            rngs = {'dropout': dropout_key}
+            rngs = {"dropout": dropout_key}
             with te.fp8_autocast(**fp8_autocast_kwargs):
-                call_kwargs = {**forward_kwargs, 'rngs': rngs}
+                call_kwargs = {**forward_kwargs, "rngs": rngs}
                 return model_apply_fn(variables, inp, **call_kwargs)
-        
+
         # Compute forward pass and get VJP function (w.r.t. variables and inp)
         output, vjp_fn = jax.vjp(forward_fn, variables, inp)
 
         # Compute gradients using VJP - returns gradients w.r.t. variables and inp
         var_grads, inp_grads = vjp_fn(grad_target)
-        
         # Return loss value and gradients
         loss_value = jnp.vdot(output, grad_target)
         return loss_value, (var_grads, inp_grads)