# coding=utf-8 # Copyright 2024 The Google Research Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Transformer LM trainer.""" import functools from clu import periodic_actions from flax.training import common_utils from flax.training import train_state import jax from jax import numpy as jnp import numpy as np import optax import ml_collections from train import model def get_state(config, net, initial_variables): """Get the train state given an experiment config, a model and initial variables. Args: config: A ConfigDict containing the configuration for the experiment. net: The model to use for training. initial_variables: The initial variables for the model. Returns: A tuple containing the train state and the learning rate schedule. """ # Learning rate schedule lr_scheduler_fn = functools.partial( lr_scheduler, learning_rate=config.learning_rate, warmup_tokens=config.warmup_tokens, final_tokens=config.max_steps, config=config, ) # Optimizer optim_fn = optax.adamw( lr_scheduler_fn, weight_decay=config.weight_decay, b1=0.9, b2=0.95 ) # Clip the gradients to prevent exploding gradients. optimizer = optax.chain(optax.clip_by_global_norm(1), optim_fn) # Initialize the train state state = train_state.TrainState.create( apply_fn=net.apply, params=initial_variables["params"], tx=optimizer ) return state, lr_scheduler_fn def lr_scheduler( n_tokens: int, learning_rate: float, warmup_tokens: int, final_tokens: int, config: ml_collections.ConfigDict, ) -> float: """Learning rate scheduler, adapted from Mikhail Grankin. The learning rate schedule is cosine decay with a warmup period. The learning rate starts at 0 and linearly increases to the given learning rate over the warmup period. After the warmup period, the learning rate decays according to a cosine schedule, with the given learning rate as the maximum value. Args: n_tokens: The number of tokens processed so far. learning_rate: The initial learning rate. warmup_tokens: The number of tokens to warm up over. final_tokens: The total number of tokens to process. config: A ConfigDict containing the configuration for the learning rate schedule. Returns: The learning rate at the given point in the schedule. """ # Decay the learning rate based on our progress. progress = (n_tokens - warmup_tokens) / max( 1, final_tokens - warmup_tokens, ) lr_mult = jnp.where( n_tokens < warmup_tokens, # Linear warmup. n_tokens / jnp.fmax(1, warmup_tokens), # Cosine learning rate decay. jnp.fmax(config.end_lr_factor, 0.5 * (1.0 + jnp.cos(np.pi * progress))), ) return learning_rate * lr_mult def get_metrics_report_progress(config, workdir, writer): """ Get the metrics for reporting progress during training. Args: config: The configuration for the experiment. workdir: The directory for storing the logs. writer: The writer object for recording the metrics. Returns: hooks: List of hooks for tracking progress. report_progress: Object for reporting progress. train_metrics: List of training metrics. """ hooks = [] # Initialize the report progress object report_progress = periodic_actions.ReportProgress( num_train_steps=config.max_steps, writer=writer) # Add metrics for profiling if the process index is 0 if jax.process_index() == 0: hooks += [report_progress, periodic_actions.Profile(logdir=workdir, num_profile_steps=5)] # Initialize the list of training metrics train_metrics = [] return hooks, report_progress, train_metrics def get_input_start_index(batch, config): inputs = jax.tree_util.tree_map(np.asarray, batch[0]) puzzles = jax.tree_util.tree_map(np.asarray, batch[1]) start_index = jax.tree_util.tree_map(np.asarray, batch[2]) levels = jax.tree_util.tree_map(np.asarray, batch[3]) cand_targets = jax.tree_util.tree_map(np.asarray, batch[4]) return inputs, puzzles, start_index, levels, cand_targets def train_one_step(p_train_step, config, state, step, dropout_rngs, train_data_iter): """ Single step of the training loop. Args: p_train_step: The training step function. config: The experiment configuration. state: The train state. step: The step number. dropout_rngs: The dropout random number generator. train_data_iter: The iterator for the train data. Returns: The updated train state and train metrics. """ with jax.profiler.StepTraceAnnotation("train", step_num=step): # Get the next batch from the iterator batch = next(train_data_iter) # Extract the inputs, start index and difficulty level from the batch inputs, _, start_index, levels, cand_targets = get_input_start_index(batch, config) # Shard across the devices inputs = common_utils.shard(jax.tree_util.tree_map(np.asarray, inputs)) start_index = common_utils.shard(jax.tree_util.tree_map(np.asarray, start_index)) levels = common_utils.shard(jax.tree_util.tree_map(np.asarray, levels)) # Explicit reshape (not common_utils.shard) so a zero-width slot dim # (K=0 no-latent baseline: cand_targets is (bs, 0, 81)) shards without # the ambiguous -1 inference. Identical to shard() when K>0. cand_targets = np.asarray(cand_targets) _nd = jax.local_device_count() cand_targets = cand_targets.reshape( (_nd, cand_targets.shape[0] // _nd) + cand_targets.shape[1:]) # Run the training step state, metrics, _ = p_train_step( state, inputs, start_index, levels, cand_targets, dropout_rng=dropout_rngs ) return state, metrics def _shard_train_batch(batch, config): """Unpack a dataset batch and shard it the way train_one_step does.""" inputs, _, start_index, levels, cand_targets = get_input_start_index( batch, config) inputs = common_utils.shard(jax.tree_util.tree_map(np.asarray, inputs)) start_index = common_utils.shard( jax.tree_util.tree_map(np.asarray, start_index)) levels = common_utils.shard(jax.tree_util.tree_map(np.asarray, levels)) cand_targets = np.asarray(cand_targets) nd = jax.local_device_count() cand_targets = cand_targets.reshape( (nd, cand_targets.shape[0] // nd) + cand_targets.shape[1:]) return inputs, start_index, levels, cand_targets def apply_grads(state, grads): """Replicated optimizer step. Used after averaging per-depth grads.""" return state.apply_gradients(grads=grads) def train_depth_mix_step(grad_step_for, p_apply, config, state, step, dropout_rngs, mix_ds, n_depths): """One optimizer step whose CE is the mean over depths 1..n_depths. Depth d is run with d latent passes on a batch whose value tokens are stage-(d-1) instances, drawn from `mix_ds` (a SudokuDataset, not tf.data). Gradients are summed and divided by n_depths, then applied once, so the update is joint rather than a serial overwrite. """ n_depths = int(max(n_depths, 1)) bs = int(config.minibatch_size) acc = None metrics_list = [] with jax.profiler.StepTraceAnnotation("train", step_num=step): for d in range(1, n_depths + 1): batch = mix_ds.sample_instance_batch(bs, d - 1) inputs, start_index, levels, cand_targets = _shard_train_batch( batch, config) grads, metrics, _ = grad_step_for(d)( state, inputs, start_index, levels, cand_targets, dropout_rng=dropout_rngs) acc = (grads if acc is None else jax.tree_util.tree_map(lambda a, b: a + b, acc, grads)) metrics_list.append(metrics) acc = jax.tree_util.tree_map(lambda a: a / n_depths, acc) state = p_apply(state, acc) def _mean_metric(key): vals = [m[key] for m in metrics_list if key in m] if not vals: return None return sum(vals) / float(len(vals)) metrics = dict(metrics_list[-1]) for key in ("loss", "ce_loss", "aux_loss", "leak_loss", "latent_sigma"): v = _mean_metric(key) if v is not None: metrics[key] = v return state, metrics def build_latent_state(inputs, start_index, levels, config, num_passes, apply_fn, rngs=None, trunc_passes=0): """Build the continuous latent thoughts (Coconut/ATC-style recurrence). Pass 0 seeds z_1 from the last-layer hidden at the last clue token; pass j reads the hidden at latent slot j to produce z_{j+1}. Slot j is active for an example only when j < k, with k = clip(level - 2, 1, K) (difficulty- matched latent budget). By default gradients flow through all passes (full BPTT). If trunc_passes=T>0, the forward chain is unchanged but thoughts older than the last T CoTs are stop-gradient'd, so the loss only backprops through those last T recurrent steps. Returns (latent_vals, slot_pos, active_full). """ num_slots = int(config.num_latent_slots) bs = inputs.shape[0] bidx = jnp.arange(bs) si3 = 3 * start_index.reshape(-1) # (bs,) slot_pos = si3[:, None] + jnp.arange(num_slots)[None, :] # (bs, K) if getattr(config, "cand_slot_mode", "level") == "depth": # Every slot the recurrence fills is active, uniformly over the batch. k = jnp.full((bs,), max(min(num_passes, num_slots), 1), dtype=jnp.int32) else: k = jnp.clip(levels.reshape(-1) - 2, 1, num_slots) # (bs,) active_full = jnp.arange(num_slots)[None, :] < k[:, None] # (bs, K) latent_vals = jnp.zeros((bs, num_slots, config.emb_dim), dtype=config.dtype) for j in range(num_passes): act_j = active_full & (jnp.arange(num_slots)[None, :] < j) _, hidden, _ = apply_fn(inputs, latent_vals, slot_pos, act_j, rngs) src = si3 - 1 + j # j=0: last clue token; j>0: latent slot j-1 z = hidden[bidx, src] # Truncated BPTT: keep the full forward chain, but detach CoTs # before the last T so earlier thoughts are treated as constants. if trunc_passes > 0 and j < num_passes - trunc_passes: z = jax.lax.stop_gradient(z) latent_vals = latent_vals.at[:, j].set(z.astype(config.dtype)) return latent_vals, slot_pos, active_full def build_latent_state_noisy(inputs, start_index, levels, config, num_passes, apply_fn, sample_slot, rng, trunc_passes=0): """One sample: noise h1, pass, noise h2, pass, ..., then readout.""" num_slots = int(config.num_latent_slots) bs = inputs.shape[0] bidx = jnp.arange(bs) si3 = 3 * start_index.reshape(-1) slot_pos = si3[:, None] + jnp.arange(num_slots)[None, :] if getattr(config, "cand_slot_mode", "level") == "depth": k = jnp.full((bs,), max(min(num_passes, num_slots), 1), dtype=jnp.int32) else: k = jnp.clip(levels.reshape(-1) - 2, 1, num_slots) active_full = jnp.arange(num_slots)[None, :] < k[:, None] latent_vals = jnp.zeros((bs, num_slots, config.emb_dim), dtype=config.dtype) keys = jax.random.split(rng, num_passes) last_sigma = jnp.zeros((), jnp.float32) for j in range(num_passes): act_j = active_full & (jnp.arange(num_slots)[None, :] < j) drop_k, noise_k = jax.random.split(keys[j]) _, hidden, _ = apply_fn( inputs, latent_vals, slot_pos, act_j, {"dropout": drop_k}) # Same read position as the deterministic recurrence and the gate: # j=0 is the last clue token, j>0 is latent slot j-1's hidden. z = hidden[bidx, si3 - 1 + j] if trunc_passes > 0 and j < num_passes - trunc_passes: z = jax.lax.stop_gradient(z) mu, log_sigma = sample_slot(z, j) log_sigma = jnp.clip(log_sigma, -12.0, 2.0) eps = jax.random.normal(noise_k, z.shape, jnp.float32) z = (mu + jnp.exp(log_sigma) * eps).astype(config.dtype) last_sigma = jnp.exp(log_sigma).mean() latent_vals = latent_vals.at[:, j].set(z) return latent_vals, slot_pos, active_full, last_sigma def train_step(state, batch, start_index, levels, cand_targets, config, hyperparams, learning_rate_fn, num_passes, dropout_rng=None, backtrack=False, inst_slot=-1, apply_update=True): """One step of the training loop. Args: state: Train state. batch: Input batch (bs, 3*81 + K) with latent placeholder slots. start_index: Number of clue cells per example. levels: Puzzle difficulty level (3..8) per example. config: Model config. hyperparams: Hyperparameter dictionary. learning_rate_fn: Learning rate function. num_passes: Number of latent recurrence passes (= curriculum stage). dropout_rng: RNG used for dropout. inst_slot: Candidate-mask slot matching the stage whose instances supply this batch's value tokens (= curriculum stage - 1). Only the leak penalty reads it. -1 derives it from num_passes. Returns: A new train state, train metrics, and computed model predictions. """ num_slots = int(config.num_latent_slots) # Extract inputs and labels from the batch inputs = batch[:, :-1] label = batch[:, 1:] # Update dropout_rng dropout_rng = jax.random.fold_in(dropout_rng, state.step) # Stochastic latents draw from their own stream. Only split when the arm is # on, so the deterministic control keeps its exact dropout sequence. n_mc = int(getattr(hyperparams, "latent_mc_samples", 0)) use_noise = bool(getattr(config, "latent_noise", False)) and n_mc > 0 noise_rng = None if use_noise: dropout_rng, noise_rng = jax.random.split(dropout_rng) mc_objective = str(getattr(hyperparams, "latent_mc_objective", "mean")) dropout_rng_dict = {"dropout": dropout_rng} def loss_fn(params): """Compute the loss function.""" net = model.TransformerLMHeadModel(config) def apply_fn(x, lv, lp, la, rngs): return net.apply({"params": params}, x, latent_values=lv, latent_positions=lp, latent_active=la, rngs=rngs) # M independent samples of (noise h1, pass, noise h2, ..., readout). # Mean CE. sigma_mean = jnp.zeros((), jnp.float32) if use_noise: def _sample_slot(z, slot): return net.apply( {"params": params}, z, slot, method=model.TransformerLMHeadModel.latent_noise_dist_slot) def _one_draw(key): build_key, read_key = jax.random.split(key) def apply_draw(x, lv, lp, la, rngs): return apply_fn(x, lv, lp, la, rngs) latent_vals, slot_pos, active_full, sig = ( build_latent_state_noisy( inputs, start_index, levels, config, num_passes, apply_draw, _sample_slot, build_key, trunc_passes=int(getattr(hyperparams, "bptt_trunc", 0)))) lg, _, cl = apply_draw( inputs, latent_vals, slot_pos, active_full, {"dropout": read_key}) return lg, cl, sig pred_logits_mc, cand_logits_mc, sig_mc = jax.vmap(_one_draw)( jax.random.split(noise_rng, n_mc)) pred_logits = pred_logits_mc.mean(axis=0) cand_logits = cand_logits_mc.mean(axis=0) sigma_mean = sig_mc.mean() else: latent_vals, slot_pos, active_full = build_latent_state( inputs, start_index, levels, config, num_passes, apply_fn, rngs=dropout_rng_dict, trunc_passes=int(getattr(hyperparams, "bptt_trunc", 0))) pred_logits, _, cand_logits = apply_fn( inputs, latent_vals, slot_pos, active_full, dropout_rng_dict) pred_logits_mc = pred_logits[None] label_one_hot = jax.nn.one_hot(label, num_classes=config.vocab_size) # The variables label_one_hot and pred_logits both are 3-dimensional tensors with # first axis corresponding to batch size, second correspondingn to sequence length # and third corresponding to the row/column/value at a particular cell assert label_one_hot.shape == pred_logits.shape, ("one hot label shape", label_one_hot.shape, label.shape, pred_logits.shape) # Calculate the cross-entropy loss along the last axis, per MC draw. # ce_loss is (M, bs, L); M is 1 on the deterministic path, so the # arithmetic below reduces to the original single-forward loss exactly. pred_logits_sol = pred_logits[:, :, :] label_one_hot_sol = label_one_hot[:, :, :] ce_loss = optax.softmax_cross_entropy( logits=pred_logits_mc, labels=label_one_hot_sol[None] ) # Apply masking to the loss: supervise only the solution region, # which now starts K latent slots after the clue block. pos = jnp.arange(label.shape[1])[None, :] # (1, L) base = 3 * jnp.asarray(start_index).reshape(-1, 1) + num_slots mask = pos >= base # Split the supervised tokens by role. The solution region is a stream of # (row, col, value) triples, so the value token of empty cell t sits at # base+3t+2; ce_loss[i] scores the token at i+1, hence the shift below. # Two of every three supervised tokens are locations, and they demand the # solver's particular linearization of a wave whose cells are actually # order-interchangeable -- pressure the promotion gate (a teacher-forced, # value-only CE(q||p)) never measures. Down-weighting them spends the # capacity on the value distribution instead. is_value = ((pos + 1 - base) % 3) == 2 loc_weight = float(getattr(hyperparams, "loc_loss_weight", 1.0)) value_weight = float(getattr(hyperparams, "value_loss_weight", 1.0)) token_w = jnp.where(is_value, value_weight, loc_weight) * mask # Per-token mean (not per-example sum): keeps the LM CE on the same # O(1) scale as the per-digit BCE below, so aux_cand_weight~1 actually # balances the two instead of the candidate signal being swamped. # The denominator stays the *unweighted* supervised count so that # re-weighting moves the location gradient without silently rescaling the # value gradient too, and weights of 1.0 reproduce the plain mean exactly. ce_denom = jnp.maximum(mask.sum(), 1.0) # Per-draw scalar first, then reduce over the draws. "mean" is the # Monte-Carlo estimate of E_z[CE]; "min" charges only the best draw, # which is what lets separate draws specialize instead of all being # pulled to the same target. per_draw_ce = (ce_loss * token_w[None]).sum(axis=(1, 2)) / ce_denom if use_noise and mc_objective == "min": avg_ce_loss = per_draw_ce.min() else: avg_ce_loss = per_draw_ce.mean() # ---- Leak penalty: probability mass placed outside the candidate set -- # The gate is CE(q||p) <= log|S|, and CE(q||p) = -log(in-set mass) # + H(q) + KL(q||p_renormalized). Measured on a stage-3 checkpoint the # leak term is what grows with curriculum depth (0.03 nats at stage 3, # 0.07 at stage 4, 0.13 at stage 5) while the within-set shape error # stays flat, so it accounts for the whole deepening deficit. Nothing # else in this objective penalizes it: the LM CE only ever rewards the # one sampled digit, and the candidate head that used to supply the set # boundary is off (aux_cand_weight=0) on the instance arm. This term # minimizes exactly the measured quantity, -log(sum of p over S). leak_weight = float(getattr(hyperparams, "leak_weight", 0.0)) if leak_weight > 0.0: slot = inst_slot if inst_slot >= 0 else num_passes - 1 slot = int(min(max(slot, 0), num_slots - 1)) # At a value position i the row/col tokens sit at i-1 and i, so the # cell this position is answering is inputs[i-1]*9 + inputs[i]. prev_tok = jnp.concatenate( [jnp.zeros((inputs.shape[0], 1), inputs.dtype), inputs[:, :-1]], axis=1) cell = jnp.clip(prev_tok.astype(jnp.int32) * 9 + inputs.astype(jnp.int32), 0, 80) bits_at = jnp.take_along_axis( cand_targets[:, slot, :].astype(jnp.int32), cell, axis=1) dv = jnp.arange(config.vocab_size) # Token ids are the digits themselves, so digit d lives at index d. in_set = ( ((bits_at[:, :, None] >> jnp.clip(dv - 1, 0, 8)[None, None, :]) & 1) > 0) & ((dv >= 1) & (dv <= 9))[None, None, :] # Shifted softmax in fp32: p <= 1 and Z >= 1, so both logs are safe # without a logsumexp import. lg = pred_logits_sol.astype(jnp.float32) pe = jnp.exp(lg - lg.max(axis=-1, keepdims=True)) z_all = pe.sum(axis=-1) z_in = (pe * in_set).sum(axis=-1) leak = jnp.log(z_all) - jnp.log(jnp.maximum(z_in, 1e-30)) # Only real empty cells: clue cells are sentinel-zeroed in the masks. leak_m = mask & is_value & (bits_at > 0) leak_loss = ((leak * leak_m).sum() / jnp.maximum(leak_m.sum(), 1.0)) else: leak_loss = jnp.zeros((), dtype=jnp.float32) # ---- Auxiliary multi-candidate BCE loss on the latent slots ---- # cand_targets: (bs, K, 81) int bitmask (bit d-1 set <=> digit d is a # candidate at that slot's stage; slot j already aligned to the puzzle's # k=level-2 budget in the data pipeline). Expand to (bs,K,81,9) multi-hot. aux_weight = float(getattr(hyperparams, "aux_cand_weight", 1.0)) if aux_weight == 0.0: # Instance arm: the superposition lives in the varying targets, not # in a set head. Drop the whole BCE graph so nothing but the LM CE # shapes the latents. zero = jnp.zeros((), dtype=pred_logits.dtype) total = avg_ce_loss + leak_weight * leak_loss return total, (pred_logits, avg_ce_loss, zero, leak_loss, sigma_mean) bits = jnp.arange(9) cand_multi_hot = ((cand_targets[..., None].astype(jnp.int32) >> bits) & 1).astype(cand_logits.dtype) # (bs,K,81,9) # Positive-weighted BCE. Candidate masks are sparse (~1-3 of 9 digits # "on"), so plain BCE collapses to predicting all-zeros. Up-weighting # the positive (candidate-present) term by pos_weight counteracts the # imbalance and forces the head to actually predict the candidate set. pos_weight = float(getattr(hyperparams, "aux_pos_weight", 5.0)) log_p = jax.nn.log_sigmoid(cand_logits) # log sigmoid(x) log_1mp = jax.nn.log_sigmoid(-cand_logits) # log(1 - sigmoid(x)) bce = -(pos_weight * cand_multi_hot * log_p + (1.0 - cand_multi_hot) * log_1mp) # (bs,K,81,9) # Only the empty cells are supervised: clue cells were sentinel-zeroed # in the data pipeline, so any cell whose target row is all-zero is a # clue and must not contribute. This makes the effective target # (#empty cells) x 9 per puzzle rather than 81 x 9. cell_mask = (cand_targets > 0).astype(cand_logits.dtype) # (bs, K, 81) # Delta weighting: consecutive stages are near-duplicates (at K=12 only # ~14% of cells change per stage, so ~97% of the target bits are copies # of the previous slot). Weighting the unchanged cells below 1 stops the # objective from being satisfied by echoing slot j-1. Slot 0 has no # predecessor, so it is fully weighted. delta_bg = float(getattr(hyperparams, "aux_delta_bg", 1.0)) if delta_bg != 1.0: changed = jnp.concatenate( [jnp.ones_like(cand_targets[:, :1], dtype=bool), cand_targets[:, 1:] != cand_targets[:, :-1]], axis=1) cell_w = cell_mask * (delta_bg + (1.0 - delta_bg) * changed.astype(cand_logits.dtype)) else: cell_w = cell_mask cell_denom = jnp.maximum(cell_w.sum(axis=2), 1e-6) # (bs, K) bce_per_cell = bce.mean(axis=3) # (bs, K, 81), mean over 9 digits bce_per_slot = ((bce_per_cell * cell_w).sum(axis=2) / cell_denom) # (bs, K), mean over empty cells if backtrack: # Strict stage-replay: supervise ONLY the stage-t readout = the last # active latent slot per example (index k-1), whose candidate target # is that stage's grid. Earlier slots and the LM CE are dropped, so # this step purely re-derives "apply f exactly t times -> stage-t # readout." Injection above still spans all active slots, so the # recurrence reaching slot k-1 is intact. k_per = active_full.sum(axis=1) # (bs,) last_idx = (k_per - 1)[:, None] # (bs, 1) slot_sel = ((jnp.arange(num_slots)[None, :] == last_idx) & active_full) # (bs, K) slot_active = slot_sel.astype(bce_per_slot.dtype) else: slot_active = active_full.astype(bce_per_slot.dtype) # (bs, K) aux_denom = jnp.maximum(slot_active.sum(), 1.0) avg_aux_loss = (bce_per_slot * slot_active).sum() / aux_denom # Backtrack (replay) steps train only from the readout (no LM CE). ce_term = 0.0 if backtrack else avg_ce_loss leak_term = 0.0 if backtrack else leak_weight * leak_loss total_loss = ce_term + aux_weight * avg_aux_loss + leak_term return total_loss, (pred_logits, avg_ce_loss, avg_aux_loss, leak_loss, sigma_mean) # Compute the learning rate and perform gradient descent step = state.step lr = learning_rate_fn(step) (loss, aux), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params) pred_logits, ce_loss, aux_loss, leak_loss, sigma_mean = aux grads = jax.lax.pmean(grads, "batch") metrics = { "step": step, "loss": loss, "learning_rate": lr, "ce_loss": ce_loss, "aux_loss": aux_loss, "leak_loss": leak_loss, # Mean of the learned per-slot sigma. This is the number that tells us # whether the noise survives training or is driven to zero. "latent_sigma": sigma_mean, "pred_logits": pred_logits, "weights": inputs.shape[0] } if apply_update: new_state = state.apply_gradients(grads=grads) return new_state, metrics, pred_logits return grads, metrics, pred_logits