Would it be correct to implement k-gradient step updates like this?

[AhmetHamzaEmra]

@jax.jit
def update_inner(model, data, other_data, z):
    def loss_fn(params):
        logits = Net().apply({'params': params}, data['kip'])
        loss = aug_cross_entropy_loss(logits=logits, labels= data['label'], z=z)
        accuracy = jnp.mean(jnp.argmax(logits, -1) == data['label'])
        return loss, accuracy

    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
    for _ in range(k_steps): # user provided parameter
        aux, grads = grad_fn(model.params)
        loss, acc = aux
        model = model.apply_gradients(grads=grads)
# Some other code later

Thanks

[marcvanzee]

Looks fine to me, assuming model is a TrainState-like object:

flax/flax/training/train_state.py

Line 22 in 33afdc9

class TrainState(struct.PyTreeNode):

[AhmetHamzaEmra]

thanks