14fromlabmlimportexperiment15fromlabml.configsimportcalculate16fromlabml\_nn.transformersimportTransformerConfigs17fromlabml\_nn.transformers.rope.experimentimportConfigsasRoPEConfigs

22classConfigs(RoPEConfigs):# , ArithmeticAutoregression):23pass

26def\_rotary\_value\_pe\_mha(c:TransformerConfigs):27fromlabml\_nn.transformers.rope.value\_peimportRotaryValuePEMultiHeadAttention28returnRotaryValuePEMultiHeadAttention(c.n\_heads,c.d\_model,1.,1.)

32calculate(TransformerConfigs.encoder\_attn,'rotary\_value',\_rotary\_value\_pe\_mha)33calculate(TransformerConfigs.decoder\_attn,'rotary\_value',\_rotary\_value\_pe\_mha)34calculate(TransformerConfigs.decoder\_mem\_attn,'rotary\_value',\_rotary\_value\_pe\_mha)

37defmain():

39experiment.create(name="rotary\_shakespeare",comment="rotary value",writers={'screen','labml'})

41conf=Configs()

43experiment.configs(conf,{

45'transformer.src\_embed':'no\_pos',46'transformer.tgt\_embed':'no\_pos',

49'transformer.encoder\_attn':'rotary\_value',

53'model':'rotary\_pe\_transformer',

56'tokenizer':'character',

58'prompt\_separator':'',

60'prompt':'It is ',

62'text':'tiny\_shakespeare',

65'seq\_len':512,

67'epochs':24,

69'batch\_size':16,

72'inner\_iterations':4,

75'd\_model':128,76'transformer.ffn.d\_ff':512,77'transformer.n\_heads':4,78'transformer.dropout':0.0,

81'optimizer.optimizer':'Adam',82'optimizer.learning\_rate':2.5e-4,8384'dataloader\_shuffle\_with\_replacement':True85})

88experiment.add\_pytorch\_models({'model':conf.model})

91withexperiment.start():

93conf.run()

97if\_\_name\_\_=='\_\_main\_\_':98main()