"tests/nn/vscode:/vscode.git/clone" did not exist on "61ece000bd1b70029270e2dccab66ffa2ca16d51"
  • Pedro Cuenca's avatar
    Flax memory efficient attention (#2889) · dc277501
    Pedro Cuenca authored
    
    
    * add use_memory_efficient params placeholder
    
    * test
    
    * add memory efficient attention jax
    
    * add memory efficient attention jax
    
    * newline
    
    * forgot dot
    
    * Rename use_memory_efficient
    
    * Keep dtype last.
    
    * Actually use key_chunk_size
    
    * Rename symbol
    
    * Apply style
    
    * Rename use_memory_efficient
    
    * Keep dtype last
    
    * Pass `use_memory_efficient_attention` in `from_pretrained`
    
    * Move JAX memory efficient attention to attention_flax.
    
    * Simple test.
    
    * style
    
    ---------
    Co-authored-by: default avatarmuhammad_hanif <muhammad_hanif@sofcograha.co.id>
    Co-authored-by: default avatarMuhHanif <48muhhanif@gmail.com>
    dc277501
unet_2d_blocks_flax.py 13.9 KB