Skip to main content

/

Docs

Nightly
v24.5
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
from max.tensor import TensorShape

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

    # You can also explicitly set the devices.
    device = cpu_device()
    new_tensor = Tensor[DType.float32, rank=2](TensorShape(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
from max.tensor import TensorShape

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

    # You can also explicitly set the devices.
    device = cpu_device()
    new_tensor = Tensor[DType.float32, rank=2](TensorShape(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 origin same as that of the tensor from which it is created.