Re: [ot][spam][crazy][data] transformer model 'attention' improvement

[k]

I'm making slow progress on running the pytorch implementations of the
paper through the same test. Attached is the start of a wrapper.
The two linked implementations have the same package name for now, so
the plan was to put both wrappers in the same file, and just check the
package to see if it has the implementation function to see which one
is installed when both wrappers are in.

import memory_efficient_attention

def moskomule_attention(queries, keys, values, query_chunk_size, key_chunk_size):


def attention(queries, keys, values, query_chunk_size, key_chunk_size):


if __name__ == '__main__':
	queries, keys, values = jax.random.normal(jax.random.PRNGKey(0), (3, 64, 8, 16))
	out = attention(queries, keys, values, query_chunk_size=4, key_chunk_size=4)
	print(out)