Python package

nn

APIs to build neural network components for deep learning models with Python.

The MAX neural network API provides two namespaces:

• max.nn: Eager-style execution.
• max.nn.legacy: Legacy graph-based API (for backward compatibility).

For functional operations like relu, softmax, and more, see the functional module.

Core API

Use these modules for all models. They provide eager-style execution with PyTorch-style syntax.

• Embedding: Vector embedding layer for token representation.
• Linear: Linear transformation layer with weights and bias.
• module: Base class for all neural network modules.
• norm: Normalization layers for training stability.
• rope: Rotary position embeddings for sequence models.
• sequential: Containers for composing modules sequentially.

Legacy API

Note

The legacy API remains available for backward compatibility. For all new models, use the max.nn API.

The legacy API provides graph-based layer implementations. See the full reference:

• legacy: Neural network legacy API documentation.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit