Mirrors_Framework/deepmind-research
1
0
Fork 0
mirror of https://github.com/google-deepmind/deepmind-research.git synced 2026-09-23 07:44:11 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Switch from jax.vmap to hk.vmap: this does not change behavior at the moment, but aligns better with API recommendations.

Browse Source 
PiperOrigin-RevId: 423789562
This commit is contained in:
Andrew Lampinen
2022-02-16 16:09:54 +00:00
committed by Diego de Las Casas
parent 766286a7a1
commit cdd3bdb636
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hierarchical_transformer_memory/hierarchical_attention/htm_attention.py
+3 -3
View File
@@ -198,7 +198,7 @@ class HierarchicalMemoryAttention(hk.Module):
key=sub_sub_top_k_contents,
value=sub_sub_top_k_contents)
return sub_attention_results
do_attention = jax.vmap(do_attention, in_axes=0)
do_attention = hk.vmap(do_attention, in_axes=0)
attention_results = do_attention(sub_inputs, top_k_contents)
attention_results = jnp.squeeze(attention_results, axis=2)
# Now collapse results across k memories
@@ -207,8 +207,8 @@ class HierarchicalMemoryAttention(hk.Module):
return attention_results
# vmap across batch
batch_within_memory_attention = jax.vmap(_within_memory_attention,
in_axes=0)
batch_within_memory_attention = hk.vmap(_within_memory_attention,
in_axes=0)
outputs = batch_within_memory_attention(
queries,
jax.lax.stop_gradient(augmented_contents),