Implements the Attention mechanism in Julia as a modular Lux layer