Question: first read memories

[pfeatherstone]

During the first run, mems == None, and the model doesn't attend to any "read" tokens, as per:

recurrent-memory-transformer-pytorch/recurrent_memory_transformer_pytorch/recurrent_memory_transformer.py

Lines 350 to 355 in 3be7d43

	if exists(read_memories):
	read_mem_length = mem_length
	read_memories = read_memories + self.read_memory_emb
	else:
	read_mem_length = 0
	read_memories = x[:, 0:0]

Why not attend to read_memory_emb, and replace with

 read_mem_length = mem_length 
 read_mems = repeat(self.read_memory_emb, 'm d -> b m d', b = batch)
 if exists(read_memories): 
     read_mems += read_memories

[lucidrains]

@pfeatherstone ohh, yea i could add it and mask it out

is that the only thing keeping it from being onnx-able?

[pfeatherstone]

Pretty much. Basically nothing can be optional with tracing and ONNX. So both mems and mask must be specified

[lucidrains]

cool, will get it done later today!

[lucidrains]

@pfeatherstone decided to take the easy way out 90de2ac let me know if this does it

[pfeatherstone]

You still need to pass mems in the forward pass when tracing. So i think the fix would be to pass mems full of zeros and make the code handle that is if it were None

[pfeatherstone]

Then probably need some control flow in a torch.jit.script function which handles the different cases

[lucidrains]

@pfeatherstone have you tried this manually?

[pfeatherstone]

Yeah I have. I could submit a PR but I didn't know if passing zeros and allowing to attend to that was ok.

[lucidrains]

@pfeatherstone it is harmless to attend to it, as the 0s get summed to the read mem positional embeddings

however, decided to add a keyword argument on forward that can mask out the read memories if need be

[lucidrains]

@pfeatherstone all good?

[pfeatherstone]

Ah yes I need to try. Can it wait till Tuesday?

[lucidrains]

@pfeatherstone yea, just checking, take your time