examples/research_projects/jax-projects/big_bird/bigbird_flax.py/Trainer/create_state
class Trainer:    def create_state(self, model, tx, num_train_steps, ckpt_dir=None):
        params = model.params
        state = TrainState.create(
            apply_fn=model.__call__,
            params=params,
            tx=tx,
            loss_fn=calculate_loss_for_nq,
        )
        if ckpt_dir is not None:
            params, opt_state, step, args, data_collator = restore_checkpoint(ckpt_dir, state)
            tx_args = {
                "lr": args.lr,
                "init_lr": args.init_lr,
                "warmup_steps": args.warmup_steps,
                "num_train_steps": num_train_steps,
                "weight_decay": args.weight_decay,
            }
            tx, lr = build_tx(**tx_args)
            state = train_state.TrainState(
                step=step,
                apply_fn=model.__call__,
                params=params,
                tx=tx,
                opt_state=opt_state,
            )
            self.args = args
            self.data_collator = data_collator
            self.scheduler_fn = lr
            model.params = params
        state = jax_utils.replicate(state)
        return state