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?
Thank you! We'll create more content like this.
Thank you for helping us improve!