IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

Identity

Identity

class max.nn.Identity

source

Bases: Module

Identity layer that passes through input unchanged.

This layer is useful for skipping certain operations (like normalization) in specific architectures such as EAGLE speculative decoding, where the draft model receives already-normalized hidden states from the target model.

When called, Identity accepts a TensorValue and returns the same tensor unchanged.

from max.driver import Accelerator, CPU, accelerator_count
from max.dtype import DType
from max.graph import DeviceRef, Graph, TensorType
from max.nn import Identity

device = Accelerator() if accelerator_count() > 0 else CPU()
device_ref = DeviceRef.from_device(device)

identity = Identity()
input_type = TensorType(DType.float32, [1, 256], device=device_ref)
with Graph("identity", input_types=[input_type]) as graph:
    output = identity(graph.inputs[0])  # output == input_tensor
    graph.output(output)