Skip to main content

Mojo package

driver

APIs to interact with devices.

Although there are several modules in this max.driver package, you'll get everything you need from this top-level driver namespace, so you don't need to import each module.

For example, the basic code you need to create tensor on CPU looks like this:

from max.driver import Tensor, cpu_device
from testing import assert_equal

def main():
tensor = Tensor[DType.float32, rank=2]((1,2))
tensor[0, 0] = 1.0

# You can also explicitly set the devices.
device = cpu_device()
new_tensor = Tensor[DType.float32, rank=2]((1,2), device)
new_tensor[0, 0] = 1.0

# You can also create slices of tensor
subtensor = tensor[:, :1]
assert_equal(subtensor[0, 0], tensor[0, 0])
from max.driver import Tensor, cpu_device
from testing import assert_equal

def main():
tensor = Tensor[DType.float32, rank=2]((1,2))
tensor[0, 0] = 1.0

# You can also explicitly set the devices.
device = cpu_device()
new_tensor = Tensor[DType.float32, rank=2]((1,2), device)
new_tensor[0, 0] = 1.0

# You can also create slices of tensor
subtensor = tensor[:, :1]
assert_equal(subtensor[0, 0], tensor[0, 0])

Modules

  • anytensor: Implements type erased generic tensor and memory types.
  • device: Defines the types and functions to interact with hardware devices.
  • device_memory:
  • tensor: Defines tensor type, which is an owned, indexible buffer allocated on a given device.
  • tensor_slice: Represents a sliced view of a tensor. The slice has lifetime same as that of the tensor from which it is created.