Skip to main content
Log in

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.

Was this page helpful?