Skip to main content

/

Docs

Nightly
v24.5
Log in

Python class

Graph

Graph​

class max.graph.Graph(name: str, forward: Callable | None = None, input_types: Iterable[Type] = (), *args, **kwargs)

Represents a single MAX graph.

A Graph is a callable routine in MAX Engine. Like functions, graphs have a name and signature. Unlike a function, which follows an imperative programming model, a Graph follows a dataflow programming model, using lazily-executed, parallel operations instead of sequential instructions.

When you instantiate a graph, you must specify the input shapes as one or more TensorType or ListType values. Then, build a sequence of ops and set the graph output with output(). For example:

from dataclasses import dataclass
import numpy as np
from max.dtype import DType
from max.graph import Graph, TensorValue, TensorType

@dataclass
class Linear:
  weight: np.ndarray
  bias: np.ndarray

  def __call__(self, x: TensorValue) -> TensorValue:
      return x @ self.weight + self.bias

linear_graph = Graph(
    "linear",
    Linear(np.ones((2, 2)), np.ones((2,))),
    input_types=[TensorType(DType.float32, (2,))],
)
from dataclasses import dataclass
import numpy as np
from max.dtype import DType
from max.graph import Graph, TensorValue, TensorType

@dataclass
class Linear:
  weight: np.ndarray
  bias: np.ndarray

  def __call__(self, x: TensorValue) -> TensorValue:
      return x @ self.weight + self.bias

linear_graph = Graph(
    "linear",
    Linear(np.ones((2, 2)), np.ones((2,))),
    input_types=[TensorType(DType.float32, (2,))],
)

You can’t call a Graph directly from Python. You must compile it and execute it with MAX Engine. For more detail, see the tutorial about how to build a graph with MAX Graph.

add_weight()​

add_weight(name: str, dtype: DType | None = None, shape: Iterable[int | str | Dim | integer] | None = None, filepath: PathLike | str | None = None, offset: int | None = None, quantization_encoding: QuantizationEncoding | None = None) → Weight

Initializes a new weight in the current graph.

  • Parameters:

    • name – The name of this weight. All weights in a graph must have unique names.
    • dtype – The DType of the weight. Defaults to Float32.
    • shape – The shape of the weight. Defaults to a scalar (shape=[1]).
    • filepath – File pointing to file containing weight value.
    • offset – Offset to weight in the file (defaults to 0).
    • quantization_encoding – Quantization encoding of the weight. Defaults to not-quantized.

  • Returns:

    The added Weight object.

  • Raises:

    • ValueError if a weight with the same name already exists in the –
    • graph. –

current​

current

inputs​

inputs*: tuple[max.graph.value.Value, ...]*

output()​

output(*outputs: Value) → None

unique_symbolic_dim()​

unique_symbolic_dim(tag: str) → SymbolicDim

Create a new symbolic dim with a different name from any other.

  • Parameters:

    tag – An additional identifier to help identify the dimension for debugging purposes.

Returns: The dimension.

weights​

weights*: dict[str, max.graph.weight.Weight]*