Skip to main content

Ctrl+K

Install

Build and Install

Usage

Quick Start Guide
Lazy Evaluation
Unified Memory
Indexing Arrays
Saving and Loading Arrays
Function Transforms
Compilation
Conversion to Other Array Frameworks
Distributed Communication
Using Streams
Exporting Functions

Examples

Linear Regression
Multi-Layer Perceptron
LLM inference
Data Parallelism
Tensor Parallelism

Ruby API Reference

Ruby API Reference

Ruby DSL

Ruby DSL
Model Declaration DSL
Builder and Graphs
Train Step DSL
Trainer Core
Trainer Data Ergonomics
Trainer Presets and Tasks
Checkpoints and Resume
Artifact Policy
Data Pipeline
Experiment DSL
Split Plan DSL
Ruby DSL Examples

C++ API Reference

Operations

Further Reading

Custom Extensions in MLX
Metal Debugger
Metal Logging
Custom Metal Kernels
Using MLX in C++

.rst

mlx.nn.MultiHeadAttention

Contents

MultiHeadAttention

mlx.nn.MultiHeadAttention#

class MultiHeadAttention(*args: Any, **kwargs: Any)#

Contents

MultiHeadAttention

By MLX Contributors

© Copyright 2023, Apple.