Skip to main content

/

Docs

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])

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?