def train(*, rng: jnp.ndarray, config: ml_collections.ConfigDict, model_cls: Type[bert_base_model.BERTBaseModel], dataset: dataset_utils.Dataset, workdir: str, writer: metric_writers.MetricWriter) -> Tuple[(train_utils.TrainState, Dict[(str, Any)], Dict[(str, Any)])]:
    'Main training loop lives in this function.\n\n  Given the model class and dataset, it prepares the items needed to run the\n  training, including the TrainState.\n\n  Args:\n    rng: Jax rng key.\n    config: Configurations of the experiment.\n    model_cls: Model class; A model has a flax_module, a loss_fn, and a\n      metrics_fn associated with it.\n    dataset: The dataset that has train_iter, eval_iter, meta_data, and\n      optionally, test_iter.\n    workdir: Directory for checkpointing.\n    writer: CLU metrics writer instance.\n\n  Returns:\n    train_state that has the state of training (including current\n      global_step, model_state, rng, and the optimizer), train_summary\n      and eval_summary which are dict of metrics. These outputs are used for\n      regression testing.\n  '
    lead_host = (jax.process_index() == 0)
    model = model_cls(config, dataset.meta_data)
    (rng, init_rng) = jax.random.split(rng)
    (params, model_state, num_trainable_params, gflops) = bert_train_utils.initialize_bert_model(model_def=model.flax_model, input_spec=dataset.meta_data['input_spec'], config=config, rngs=init_rng)
    optimizer = jax.jit(optimizers.get_optimizer(config).create, backend='cpu')(params)
    (rng, train_rng) = jax.random.split(rng)
    train_state = train_utils.TrainState(global_step=0, optimizer=optimizer, model_state=model_state, rng=train_rng, accum_train_time=0)
    start_step = train_state.global_step
    if config.checkpoint:
        (train_state, start_step) = train_utils.restore_checkpoint(workdir, train_state)
    if ((start_step == 0) and (config.get('init_from') is not None)):
        restored_model_cfg = config.init_from.get('model_config')
        init_checkpoint_path = config.init_from.get('checkpoint_path')
        restored_train_state = pretrain_utils.restore_pretrained_checkpoint(init_checkpoint_path, train_state, assert_exist=True)
        train_state = model.init_from_train_state(train_state, restored_train_state, restored_model_cfg)
        del restored_train_state
    train_state = jax_utils.replicate(train_state)
    del params
    (total_steps, steps_per_epoch) = train_utils.get_num_training_steps(config, dataset.meta_data)
    learning_rate_fn = lr_schedules.get_learning_rate_fn(config)
    train_step_pmapped = jax.pmap(functools.partial(train_step, flax_model=model.flax_model, learning_rate_fn=learning_rate_fn, loss_fn=model.loss_function, metrics_fn=model.get_metrics_fn('train'), config=config, debug=config.debug_train), axis_name='batch', donate_argnums=(0, 1))
    eval_step_pmapped = jax.pmap(functools.partial(eval_step, flax_model=model.flax_model, metrics_fn=model.get_metrics_fn('validation'), all_gather=config.get('global_metrics', False), debug=config.debug_eval), axis_name='batch', donate_argnums=(1,))
    if ('fewshot' in config):
        representation_fn_pmaped = jax.pmap(functools.partial(representation_fn, flax_model=model.flax_model, representation_layer=config.fewshot.representation_layer), donate_argnums=(1,), axis_name='batch')
        fewshotter = bert_train_utils.BERTFewShotEvaluator(representation_fn_pmaped, config.fewshot)
    log_eval_steps = (config.get('log_eval_steps') or steps_per_epoch)
    if (not log_eval_steps):
        raise ValueError("'log_eval_steps' should be specified in the config.")
    checkpoint_steps = (config.get('checkpoint_steps') or log_eval_steps)
    log_summary_steps = (config.get('log_summary_steps') or log_eval_steps)
    total_eval_steps = int(np.ceil((dataset.meta_data['num_eval_examples'] / config.batch_size)))
    steps_per_eval = (config.get('steps_per_eval') or total_eval_steps)
    compute_global_metrics = False
    if (config.get('global_metrics', False) and lead_host):
        compute_global_metrics = True
    if compute_global_metrics:
        global_metrics_evaluator = bert_train_utils.BERTGlobalEvaluator(config.global_metrics)
    (train_metrics, extra_training_logs) = ([], [])
    (train_summary, eval_summary) = (None, None)
    chrono = train_utils.Chrono(first_step=start_step, total_steps=total_steps, steps_per_epoch=steps_per_epoch, global_bs=config.batch_size, accum_train_time=int(jax_utils.unreplicate(train_state.accum_train_time)), example_type='example')
    logging.info('Starting training loop at step %d.', (start_step + 1))
    report_progress = periodic_actions.ReportProgress(num_train_steps=total_steps, writer=writer)
    hooks = [report_progress]
    if (config.get('xprof', True) and lead_host):
        hooks.append(periodic_actions.Profile(num_profile_steps=5, logdir=workdir))
    if (start_step == 0):
        step0_log = {'num_trainable_params': num_trainable_params}
        if gflops:
            step0_log['gflops'] = gflops
        writer.write_scalars(1, step0_log)
    for step in range((start_step + 1), (total_steps + 1)):
        with jax.profiler.StepTraceContext('train', step_num=step):
            train_batch = next(dataset.train_iter)
            (train_state, t_metrics, lr) = train_step_pmapped(train_state=train_state, batch=train_batch)
            train_metrics.append(t_metrics)
            extra_training_logs.append({'learning_rate': lr})
        for h in hooks:
            h(step)
        chrono.pause()
        if (((step % log_summary_steps) == 1) or (step == total_steps)):
            if lead_host:
                chrono.tick(step, writer=writer)
            train_summary = train_utils.log_train_summary(step=step, train_metrics=jax.tree_map(train_utils.unreplicate_and_get, train_metrics), extra_training_logs=jax.tree_map(train_utils.unreplicate_and_get, extra_training_logs), writer=writer)
            (train_metrics, extra_training_logs) = ([], [])
            if (((step % log_eval_steps) == 1) or (step == total_steps)):
                with report_progress.timed('eval'):
                    eval_metrics = []
                    train_state = train_utils.sync_model_state_across_replicas(train_state)
                    for _ in range(steps_per_eval):
                        eval_batch = next(dataset.valid_iter)
                        (e_metrics, e_output, e_batch) = eval_step_pmapped(train_state=train_state, batch=eval_batch)
                        eval_metrics.append(train_utils.unreplicate_and_get(e_metrics))
                        if compute_global_metrics:
                            e_batch_mask = train_utils.unreplicate_and_get(e_batch['batch_mask']).astype(bool)
                            t_key = ('label' if ('label' in e_batch) else 'targets')
                            global_metrics_evaluator.add_batch_of_examples(target=train_utils.unreplicate_and_get(e_batch[t_key])[e_batch_mask], output=train_utils.unreplicate_and_get(e_output)[e_batch_mask])
                            del e_batch, e_output, e_batch_mask
                    eval_global_metrics_summary = None
                    if compute_global_metrics:
                        if (len(global_metrics_evaluator) != dataset.meta_data['num_eval_examples']):
                            raise ValueError(f"Number of eval examples should be {dataset.meta_data['num_eval_examples']}, but it is {len(global_metrics_evaluator)}.")
                        eval_global_metrics_summary = global_metrics_evaluator.compute_metrics(clear_annotations=True)
                    eval_summary = train_utils.log_eval_summary(step=step, eval_metrics=eval_metrics, extra_eval_summary=eval_global_metrics_summary, writer=writer)
                writer.flush()
                del eval_metrics, eval_global_metrics_summary
        if (((((step % checkpoint_steps) == 0) and (step > 0)) or (step == total_steps)) and config.checkpoint):
            with report_progress.timed('checkpoint'):
                train_state = train_utils.sync_model_state_across_replicas(train_state)
                if lead_host:
                    train_state.replace(accum_train_time=chrono.accum_train_time)
                    train_utils.save_checkpoint(workdir, train_state)
        if ('fewshot' in config):
            if (((step % config.fewshot.log_eval_steps) == 1) or (step == total_steps)):
                with report_progress.timed('fewshot'):
                    results = fewshotter.run_all(train_state, config.fewshot.datasets)
                    fewshotter.log_fewshot_summary(writer=writer, step=step, results=results)
                    del results
                    writer.write_scalars(step, {'zz/epoch': (step / steps_per_epoch)})
                writer.flush()
        chrono.resume()
    jax.random.normal(jax.random.PRNGKey(0), ()).block_until_ready()
    return (train_state, train_summary, eval_summary)
Main training loop lives in this function.

Given the model class and dataset, it prepares the items needed to run the
training, including the TrainState.

Args:
  rng: Jax rng key.
  config: Configurations of the experiment.
  model_cls: Model class; A model has a flax_module, a loss_fn, and a
    metrics_fn associated with it.
  dataset: The dataset that has train_iter, eval_iter, meta_data, and
    optionally, test_iter.
  workdir: Directory for checkpointing.
  writer: CLU metrics writer instance.

Returns:
  train_state that has the state of training (including current
    global_step, model_state, rng, and the optimizer), train_summary
    and eval_summary which are dict of metrics. These outputs are used for
    regression testing.
def train(*, rng: jnp.ndarray, config: ml_collections.ConfigDict, model_cls: Type[bert_base_model.BERTBaseModel], dataset: dataset_utils.Dataset, workdir: str, writer: metric_writers.MetricWriter) -> Tuple[(train_utils.TrainState, Dict[(str, Any)], Dict[(str, Any)])]:
    'Main training loop lives in this function.\n\n  Given the model class and dataset, it prepares the items needed to run the\n  training, including the TrainState.\n\n  Args:\n    rng: Jax rng key.\n    config: Configurations of the experiment.\n    model_cls: Model class; A model has a flax_module, a loss_fn, and a\n      metrics_fn associated with it.\n    dataset: The dataset that has train_iter, eval_iter, meta_data, and\n      optionally, test_iter.\n    workdir: Directory for checkpointing.\n    writer: CLU metrics writer instance.\n\n  Returns:\n    train_state that has the state of training (including current\n      global_step, model_state, rng, and the optimizer), train_summary\n      and eval_summary which are dict of metrics. These outputs are used for\n      regression testing.\n  '
    lead_host = (jax.process_index() == 0)
    model = model_cls(config, dataset.meta_data)
    (rng, init_rng) = jax.random.split(rng)
    (params, model_state, num_trainable_params, gflops) = bert_train_utils.initialize_bert_model(model_def=model.flax_model, input_spec=dataset.meta_data['input_spec'], config=config, rngs=init_rng)
    optimizer = jax.jit(optimizers.get_optimizer(config).create, backend='cpu')(params)
    (rng, train_rng) = jax.random.split(rng)
    train_state = train_utils.TrainState(global_step=0, optimizer=optimizer, model_state=model_state, rng=train_rng, accum_train_time=0)
    start_step = train_state.global_step
    if config.checkpoint:
        (train_state, start_step) = train_utils.restore_checkpoint(workdir, train_state)
    if ((start_step == 0) and (config.get('init_from') is not None)):
        restored_model_cfg = config.init_from.get('model_config')
        init_checkpoint_path = config.init_from.get('checkpoint_path')
        restored_train_state = pretrain_utils.restore_pretrained_checkpoint(init_checkpoint_path, train_state, assert_exist=True)
        train_state = model.init_from_train_state(train_state, restored_train_state, restored_model_cfg)
        del restored_train_state
    train_state = jax_utils.replicate(train_state)
    del params
    (total_steps, steps_per_epoch) = train_utils.get_num_training_steps(config, dataset.meta_data)
    learning_rate_fn = lr_schedules.get_learning_rate_fn(config)
    train_step_pmapped = jax.pmap(functools.partial(train_step, flax_model=model.flax_model, learning_rate_fn=learning_rate_fn, loss_fn=model.loss_function, metrics_fn=model.get_metrics_fn('train'), config=config, debug=config.debug_train), axis_name='batch', donate_argnums=(0, 1))
    eval_step_pmapped = jax.pmap(functools.partial(eval_step, flax_model=model.flax_model, metrics_fn=model.get_metrics_fn('validation'), all_gather=config.get('global_metrics', False), debug=config.debug_eval), axis_name='batch', donate_argnums=(1,))
    if ('fewshot' in config):
        representation_fn_pmaped = jax.pmap(functools.partial(representation_fn, flax_model=model.flax_model, representation_layer=config.fewshot.representation_layer), donate_argnums=(1,), axis_name='batch')
        fewshotter = bert_train_utils.BERTFewShotEvaluator(representation_fn_pmaped, config.fewshot)
    log_eval_steps = (config.get('log_eval_steps') or steps_per_epoch)
    if (not log_eval_steps):
        raise ValueError("'log_eval_steps' should be specified in the config.")
    checkpoint_steps = (config.get('checkpoint_steps') or log_eval_steps)
    log_summary_steps = (config.get('log_summary_steps') or log_eval_steps)
    total_eval_steps = int(np.ceil((dataset.meta_data['num_eval_examples'] / config.batch_size)))
    steps_per_eval = (config.get('steps_per_eval') or total_eval_steps)
    compute_global_metrics = False
    if (config.get('global_metrics', False) and lead_host):
        compute_global_metrics = True
    if compute_global_metrics:
        global_metrics_evaluator = bert_train_utils.BERTGlobalEvaluator(config.global_metrics)
    (train_metrics, extra_training_logs) = ([], [])
    (train_summary, eval_summary) = (None, None)
    chrono = train_utils.Chrono(first_step=start_step, total_steps=total_steps, steps_per_epoch=steps_per_epoch, global_bs=config.batch_size, accum_train_time=int(jax_utils.unreplicate(train_state.accum_train_time)), example_type='example')
    logging.info('Starting training loop at step %d.', (start_step + 1))
    report_progress = periodic_actions.ReportProgress(num_train_steps=total_steps, writer=writer)
    hooks = [report_progress]
    if (config.get('xprof', True) and lead_host):
        hooks.append(periodic_actions.Profile(num_profile_steps=5, logdir=workdir))
    if (start_step == 0):
        step0_log = {'num_trainable_params': num_trainable_params}
        if gflops:
            step0_log['gflops'] = gflops
        writer.write_scalars(1, step0_log)
    for step in range((start_step + 1), (total_steps + 1)):
        with jax.profiler.StepTraceContext('train', step_num=step):
            train_batch = next(dataset.train_iter)
            (train_state, t_metrics, lr) = train_step_pmapped(train_state=train_state, batch=train_batch)
            train_metrics.append(t_metrics)
            extra_training_logs.append({'learning_rate': lr})
        for h in hooks:
            h(step)
        chrono.pause()
        if (((step % log_summary_steps) == 1) or (step == total_steps)):
            if lead_host:
                chrono.tick(step, writer=writer)
            train_summary = train_utils.log_train_summary(step=step, train_metrics=jax.tree_map(train_utils.unreplicate_and_get, train_metrics), extra_training_logs=jax.tree_map(train_utils.unreplicate_and_get, extra_training_logs), writer=writer)
            (train_metrics, extra_training_logs) = ([], [])
            if (((step % log_eval_steps) == 1) or (step == total_steps)):
                with report_progress.timed('eval'):
                    eval_metrics = []
                    train_state = train_utils.sync_model_state_across_replicas(train_state)
                    for _ in range(steps_per_eval):
                        eval_batch = next(dataset.valid_iter)
                        (e_metrics, e_output, e_batch) = eval_step_pmapped(train_state=train_state, batch=eval_batch)
                        eval_metrics.append(train_utils.unreplicate_and_get(e_metrics))
                        if compute_global_metrics:
                            e_batch_mask = train_utils.unreplicate_and_get(e_batch['batch_mask']).astype(bool)
                            t_key = ('label' if ('label' in e_batch) else 'targets')
                            global_metrics_evaluator.add_batch_of_examples(target=train_utils.unreplicate_and_get(e_batch[t_key])[e_batch_mask], output=train_utils.unreplicate_and_get(e_output)[e_batch_mask])
                            del e_batch, e_output, e_batch_mask
                    eval_global_metrics_summary = None
                    if compute_global_metrics:
                        if (len(global_metrics_evaluator) != dataset.meta_data['num_eval_examples']):
                            raise ValueError(f"Number of eval examples should be {dataset.meta_data['num_eval_examples']}, but it is {len(global_metrics_evaluator)}.")
                        eval_global_metrics_summary = global_metrics_evaluator.compute_metrics(clear_annotations=True)
                    eval_summary = train_utils.log_eval_summary(step=step, eval_metrics=eval_metrics, extra_eval_summary=eval_global_metrics_summary, writer=writer)
                writer.flush()
                del eval_metrics, eval_global_metrics_summary
        if (((((step % checkpoint_steps) == 0) and (step > 0)) or (step == total_steps)) and config.checkpoint):
            with report_progress.timed('checkpoint'):
                train_state = train_utils.sync_model_state_across_replicas(train_state)
                if lead_host:
                    train_state.replace(accum_train_time=chrono.accum_train_time)
                    train_utils.save_checkpoint(workdir, train_state)
        if ('fewshot' in config):
            if (((step % config.fewshot.log_eval_steps) == 1) or (step == total_steps)):
                with report_progress.timed('fewshot'):
                    results = fewshotter.run_all(train_state, config.fewshot.datasets)
                    fewshotter.log_fewshot_summary(writer=writer, step=step, results=results)
                    del results
                    writer.write_scalars(step, {'zz/epoch': (step / steps_per_epoch)})
                writer.flush()
        chrono.resume()
    jax.random.normal(jax.random.PRNGKey(0), ()).block_until_ready()
    return (train_state, train_summary, eval_summary)