Calling helper function with LayoutTensor value

Mojo
1

I have a function foo that is part of a kernel. It’s called with a LayoutTensor. I want to use a separate helper function bar that works on scalars coming from the LayoutTensor. How do I do that?

from layout.layout_tensor import (
    Layout,
    LayoutTensor,
)
from layout.layout import UNKNOWN_VALUE
from layout.tensor_builder import LayoutTensorBuild as tb
from layout.tensor_builder import dynamic, static

fn bar[dtype: DType](one: Scalar[dtype]):
    print(one)

fn foo[dtype: DType, corners_layout: Layout](
    corners: LayoutTensor[
        dtype, corners_layout, MutableAnyOrigin
    ],  # x1, y1, x2, y2
):
    bar(corners[0,0])

def main():
    alias N = 16

    alias corners_layout = Layout.row_major(N, 4)
    var t_2d_static = tb[DType.float32]().row_major(
        static[16](), static[4]()
    ).alloc()
    foo[DType.float32, corners_layout](t_2d_static)

The code above fails with

/Users/mseritan/dev/hackathon-may-25/nms.mojo/crash.mojo:17:8: error: invalid call to 'bar': argument #0 cannot be converted from 'SIMD[dtype, __init__[::Origin[::Bool(IntTuple(1), IntTuple(1)).size()]' to 'SIMD[dtype, 1]'
    bar(corners[0,0])
    ~~~^~~~~~~~~~~~~~
/Users/mseritan/dev/hackathon-may-25/nms.mojo/crash.mojo:17:16: note: types parameters include unfolded expression at parser time; try rebinding to a consistent type?
    bar(corners[0,0])
               ^
/Users/mseritan/dev/hackathon-may-25/nms.mojo/crash.mojo:9:4: note: function declared here
fn bar[dtype: DType](one: Scalar[dtype]):
   ^
mojo: error: failed to parse the provided Mojo source module
2

Looks like that’s missing a rebind (with nightly mojo 25.4.0.dev2025051405 (8af0c6dd))

fn foo[dtype: DType, corners_layout: Layout](
    corners: LayoutTensor[
        dtype, corners_layout, MutableAnyOrigin
    ],  # x1, y1, x2, y2
):
    bar(rebind[Scalar[dtype]](corners[0, 0]))
3

Thanks a lot @Ehsan , this works!

I have a follow up question. Since I have 4 parameters, corners[0,0], corners[0,1] the rebind is a bit verbose. Is there a better Mojo construct here? Can I pass a slice of the TensorLayout?

4

This topic was automatically closed 7 days after the last reply. New replies are no longer allowed.