code: def update_hparams_for_tpu(hparams):
    
  
    # Adafactor uses less memory than Adam.
    # switch to Adafactor with its recommended learning rate scheme.
    hparams.optimizer = "Adafactor"
    hparams.learning_rate_schedule = "rsqrt_decay"
    hparams.learning_rate_warmup_steps = 10000
  
    # Avoid an expensive concat on TPU.
    # >1 shards helps with faster parameter distribution on multi-GPU machines
    hparams.symbol_modality_num_shards = 1
  
    # Adaptive batch sizes and sequence lengths are not supported on TPU.
    # Instead, every batch has the same sequence length and the same batch size.
    # Longer sequences are dropped and shorter ones are padded.
    #
    # It is therefore suggested to use a problem where examples have been combined
    # to a longer length, e.g. the "_packed" problems.
    #
    # For problems with variable sequence lengths, this parameter controls the
    # maximum sequence length.  Shorter sequences are dropped and longer ones
    # are padded.
    #
    # For problems with fixed sequence lengths - e.g. the "_packed" problems,
    # this hyperparameter is ignored.
    hparams.max_length = 64
  
    # TPUs have less memory than GPUs, so decrease the batch size
    hparams.batch_size = 2048
  
    # Using noise broadcast in the dropout layers saves memory during training.
    hparams.attention_dropout_broadcast_dims = "0,1"  # batch, heads
    hparams.relu_dropout_broadcast_dims = "1"  # length
    hparams.layer_prepostprocess_dropout_broadcast_dims = "1"  # length
    return hparams
docs: Change hparams to be compatible with TPU training.