code: def imagetransformer_b12l_4h_b256_uncond_dr03_rel_tpu():
    
    hparams = imagetransformer_b12l_4h_b256_uncond_dr03_tpu()
    hparams.shared_rel = True
    hparams.dec_attention_type = cia.AttentionType.RELATIVE_LOCAL_1D
    return hparams
docs: works very well on 4x4.