INNER CODE UNIT · Go

testMultiHeadAttention

takara-ai/go-attention · cmd/demo/main.go:50

func testMultiHeadAttention() {
    // Configure multi-head attention
    config := attention.MultiHeadConfig{
        NumHeads:    4,
        DModel:      64,
        DKey:        16,  // DModel / NumHeads
        DValue:      16,  // DModel / NumHeads
        DropoutRate: 0.1,
    }

    // Create multi-head attention module
    mha, err := attention.NewMultiHeadAttention(config)
    if err != nil {
        log.Fatal(err)
    }

    // Create sample input with dimensions:
    // - batch_size: number of sequences to process in parallel

View source record →