Python module
type
Library for graph value types.
BufferType
class max.graph.type.BufferType(dtype, shape, device)
A symbolic buffer type.
This is a reference to a tensor that can be mutated in place.
Constructs a tensor type.
-
Parameters:
as_tensor()
as_tensor()
Returns the analogous tensor type.
-
Return type:
from_mlir()
classmethod from_mlir(type)
Constructs a buffer type from an MLIR type.
-
Parameters:
-
- t – The MLIR Type object to parse into a buffer type.
- type (
BufferType)
-
Returns:
-
The buffer type represented by the MLIR Type value.
-
Return type:
to_mlir()
to_mlir()
Converts to an mlir.Type instance.
-
Returns:
-
An
mlir.Typein the specified Context. -
Return type:
-
BufferType
DeviceKind
class max.graph.type.DeviceKind(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)
A device type representation.
CPU
CPU = 'cpu'
GPU
GPU = 'gpu'
from_string()
static from_string(txt)
-
Return type:
DeviceRef
class max.graph.type.DeviceRef(device_type, id=0)
A symbolic device representation.
DeviceRef type representation consists of a DeviceKind and an id. This is a direct representation of the device attribute in mlir.
The following example demonstrates how to create and use device references:
from max.graph import DeviceRef
gpu_device = DeviceRef.GPU()
print(gpu_device) # Outputs: gpu:0
# Create a CPU device with specific id
cpu_device = DeviceRef.CPU(id=1)
print(cpu_device) # Outputs: cpu:1
from max.graph import DeviceRef
gpu_device = DeviceRef.GPU()
print(gpu_device) # Outputs: gpu:0
# Create a CPU device with specific id
cpu_device = DeviceRef.CPU(id=1)
print(cpu_device) # Outputs: cpu:1
-
Parameters:
-
- device_type (
DeviceKind) - id (
int)
- device_type (
CPU()
static CPU(id=0)
Static Method for creating a CPU device.
GPU()
static GPU(id=0)
Static Method for creating a GPU device.
device_type
device_type: DeviceKind
from_device()
static from_device(device)
from_mlir()
static from_mlir(attr)
Returns a device from mlir attribute
-
Parameters:
-
attr (
DeviceRefAttr) -
Return type:
id
id: int
is_cpu()
is_cpu()
Returns true if the device is a CPU device.
-
Return type:
is_gpu()
is_gpu()
Returns true if the device is a GPU device.
-
Return type:
to_device()
to_device()
Convert device reference to a concrete driver Device.
-
Return type:
to_mlir()
to_mlir()
Returns a mlir attribute representing device.
-
Return type:
-
DeviceRefAttr
TensorType
class max.graph.type.TensorType(dtype, shape, device)
A symbolic TensorType.
This is not an eager tensor type! This contains no actual data, but instead represents the type of a value at some point in time during model execution.
Most internal values in a model will be tensors. This type represents
their element type (dtype) and dimensions (dims) at a specific point during
model computation. It allows us to do some optimistic optimizations and
shape inference during graph construction, and to provide more detailed
shape information to the compiler for further optimization passes.
The following example shows how to create a tensor type with static dimensions and access its properties:
from max.graph import TensorType
from max.dtype import DType
# Create a tensor type with float32 elements and static dimensions 2x3
tensor_type = TensorType(DType.float32, (2, 3))
print(tensor_type.dtype) # Outputs: DType.float32
print(tensor_type.shape) # Outputs: [2, 3]
from max.graph import TensorType
from max.dtype import DType
# Create a tensor type with float32 elements and static dimensions 2x3
tensor_type = TensorType(DType.float32, (2, 3))
print(tensor_type.dtype) # Outputs: DType.float32
print(tensor_type.shape) # Outputs: [2, 3]
It can also represent a fully dynamic rank tensor. The presence of dynamic rank tensors in a graph will often degrade performance dramatically and prevents many classes of optimizations.
An optional device (device) can also be provided to indicate the explicit
device the tensor is associated with.
Constructs a tensor type.
-
Parameters:
as_buffer()
as_buffer()
Returns the analogous buffer type.
-
Return type:
from_mlir()
classmethod from_mlir(type)
Constructs a tensor type from an MLIR type.
-
Parameters:
-
- t – The MLIR Type object to parse into a tensor type.
- type (
TensorType)
-
Returns:
-
The tensor type represented by the MLIR Type value.
-
Return type:
to_mlir()
to_mlir()
Converts to an mlir.Type instance.
-
Returns:
-
An
mlir.Typein the specified Context. -
Return type:
-
TensorType
Type
class max.graph.type.Type
Represents any possible type for Graph values.
Every Value in the Graph has a Type, and that type is represented by an Type. This type may be inspected to get finer-grained types and learn more about an individual Value.
The following example shows how to work with types in a graph:
from max.graph import Graph, TensorType
from max.dtype import DType
with Graph() as g:
# Create a tensor constant with a specific type
tensor_type = TensorType(DType.float32, [2, 3])
# The type can be inspected to get information about the value
print(f"Tensor element type: {tensor_type.dtype}") # Outputs: DType.float32
print(f"Tensor shape: {tensor_type.shape}") # Outputs: [2, 3]
from max.graph import Graph, TensorType
from max.dtype import DType
with Graph() as g:
# Create a tensor constant with a specific type
tensor_type = TensorType(DType.float32, [2, 3])
# The type can be inspected to get information about the value
print(f"Tensor element type: {tensor_type.dtype}") # Outputs: DType.float32
print(f"Tensor shape: {tensor_type.shape}") # Outputs: [2, 3]
from_mlir()
static from_mlir(t)
Constructs a type from an MLIR type.
-
Parameters:
-
t (
MlirType) – The MLIR Type object to parse into a type. -
Returns:
-
The type represented by the MLIR Type value.
-
Return type:
to_mlir()
to_mlir()
Converts to an mlir.Type instance.
-
Returns:
-
An
mlir.Typein the specified Context. -
Return type:
-
MlirType
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!