The moskomule attention implementation is matching for me if I perform
the square root myself outside it. The project does not appear to
install properly, the test and the library file run adjacent in the
same folder.
import jax, torch, math
import memory_efficient_attention
def moskomule_attention(queries, keys, values, query_chunk_size, key_chunk_size):
queries /= math.sqrt(keys.shape[-1])
return memory_efficient_attention.efficient_attention(queries, keys, values, chunk_size=key_chunk_size, checkpointing = True, out_of_place = False)
def attention(queries, keys, values, query_chunk_size, key_chunk_size):
queries = queries.permute(0, 2, 1, 3)
keys = keys.permute(0, 2, 1, 3)
values = values.permute(0, 2, 1, 3)
if hasattr(memory_efficient_attention, 'efficient_attention'):
return moskomule_attention(queries, keys, values, query_chunk_size, key_chunk_size)
if __name__ == '__main__':
queries, keys, values = torch.from_numpy(jax.random.normal(jax.random.PRNGKey(0), (3, 1, 64, 8, 16)).to_py())
out = attention(queries, keys, values, query_chunk_size=4, key_chunk_size=4)
print(out)
