from layout.tensor_builder import LayoutTensorBuild as tensor_builder
from layout.layout_tensor import (
Layout,
LayoutTensor,
)
fn example_tensor[rows: Int, columns: Int]() -> LayoutTensor[DType.int32, Layout.row_major(rows, columns), MutableAnyOrigin]:
var t = tensor_builder[DType.int32]().row_major[rows, columns]().alloc()
@parameter
for r in range(rows):
@parameter
for c in range(columns):
t[r, c] = (r + 1) * 100 + c
return t
fn main():
C = example_tensor[8, 16]
when I run it:
$ mojo test.mojo
test.mojo:14:12: error: cannot implicitly convert 'LayoutTensor[int32, row_major[::Origin[::Bool(_to_int_tuple[::VariadicList[::Int]]()), MutableAnyOrigin]' value to 'LayoutTensor[int32, row_major(rows, columns), MutableAnyOrigin]'
return t
Summoning @lukas who knows how this is supposed to work..
I suspect this is a case where the Mojo compiler can’t prove that the return type of the builder.row_major() is the same as the type of the return value, but the types are the same. This is something that we’d need to fix either in the compiler, or in the definition of the builder struct.
As a workaround, I think you can use a rebind. The following works for me:
Here is the complete example, it compiles and runs fine using Mojo nightly.
from layout.tensor_builder import LayoutTensorBuild as tensor_builder
from layout.layout_tensor import (
Layout,
LayoutTensor,
)
fn example_tensor[rows: Int, columns: Int]() -> LayoutTensor[DType.int32, Layout.row_major(rows, columns), MutableAnyOrigin]:
var t = tensor_builder[DType.int32]().row_major[rows, columns]().alloc()
@parameter
for r in range(rows):
@parameter
for c in range(columns):
t[r, c] = (r + 1) * 100 + c
return rebind[LayoutTensor[DType.int32, Layout.row_major(rows, columns), MutableAnyOrigin]](t)
fn main():
var C = example_tensor[8, 16]()
print(C)
Can you post your example (all lines) so that we can figure out what is different?
After upgrading Mojo, it can run. But the code does not generate the expected matrix.
Now I have found a way to get correct results, but not through type rebinding.
Here is the full code:
from layout.tensor_builder import LayoutTensorBuild as tensor_builder, static
from layout.layout_tensor import (
Layout,
LayoutTensor,
)
fn example_tensor_incorrect[rows: Int, columns: Int]() -> LayoutTensor[DType.int32, Layout.row_major(rows, columns), MutableAnyOrigin]:
var t = tensor_builder[DType.int32]().row_major[rows, columns]().alloc()
@parameter
for r in range(rows):
@parameter
for c in range(columns):
t[r, c] = (r + 1) * 100 + c
return rebind[LayoutTensor[DType.int32, Layout.row_major(rows, columns), MutableAnyOrigin]](t)
fn example_tensor_correct[rows: Int, columns: Int]() -> LayoutTensor[
DType.int32,
Layout.row_major(rows, columns),
MutableAnyOrigin,
]:
var t = tensor_builder[DType.int32]().row_major( static[rows](), static[columns]()
).alloc()
for r in range(rows):
for c in range(columns): t[r, c] = (r + 1) * 100 + c
return t
fn main():
C = example_tensor_incorrect[4, 5]()
print(C)
print('-' * 30)
C = example_tensor_correct[4, 5]()
print(C)