Modular•7mo ago

Assigning slice of a tensor to a variable?

As best I can tell there is no direct way to assign a slice of a tensor to a new variable. I am looking for something Python-esque like

fn blah():
    var t = Tensor[DType.float64](
        (3, 1),
        -1,
        0,
        1,
    )
    var x = t[Index(0:2)]
    print(x)

fn blah():
    var t = Tensor[DType.float64](
        (3, 1),
        -1,
        0,
        1,
    )
    var x = t[Index(0:2)]
    print(x)

But even the official docs show for loops iterating over a Tensor to populate a new array. Is there a better way to achieve this? It becomes more of an issue with a larger Tensor where I want to assign multiple variables to different slices.

4 Replies

Jack Clayton•7mo ago

We haven't implemented slicing in Mojo Tensors yet like you can with a String for example:

fn main():
    var s = String("mojo")
    var x = s[0:2]
    print(x)

fn main():
    var s = String("mojo")
    var x = s[0:2]
    print(x)

Could you raise a feature request for that? https://github.com/modularml/mojo/issues/new/choose Also we're talking about what we're going to do with Tensor at the moment, there are some performance issues just so you're aware.

DJ_HalfcourtViolationOP•7mo ago

Will do, thanks In the meantime how would you suggest getting slices for multiple variables? Something like

fn blah():
    var t = Tensor[DType.float64](
        (3, 1),
        -1,
        0,
        1,
        0
    )
    var x: Tensor[DType.float64] = Tensor[DType.float64](2, 1)
    var y: Tensor[DType.float64] = Tensor[DType.float64](2, 1)
    for i in range(4):
        if i < 2:
            x[i] = t[i, 0]
        else:
            y[i] = t[i, 0]    
    print(x)

fn blah():
    var t = Tensor[DType.float64](
        (3, 1),
        -1,
        0,
        1,
        0
    )
    var x: Tensor[DType.float64] = Tensor[DType.float64](2, 1)
    var y: Tensor[DType.float64] = Tensor[DType.float64](2, 1)
    for i in range(4):
        if i < 2:
            x[i] = t[i, 0]
        else:
            y[i] = t[i, 0]    
    print(x)

Jack Clayton•7mo ago

Yeah say if you wanted your original slice you could do:

from tensor import Tensor

fn main():
    var t = Tensor[DType.float64](
        (4, 1),
        -1,
        0,
        1,
        0
    )
    var x = Tensor[DType.float64](2)
    x.store(0, t.load[width=2]())
    print(x)

from tensor import Tensor

fn main():
    var t = Tensor[DType.float64](
        (4, 1),
        -1,
        0,
        1,
        0
    )
    var x = Tensor[DType.float64](2)
    x.store(0, t.load[width=2]())
    print(x)

mad alex 1997•7mo ago

ND slicing is hard, it was one of the things I could not get quite right when I was trying to make an ND-Array type from scratch unless I manually built overloads for the getting and setting dunder methods for the struct for each possible dim of the array. This is why numpy had (and maybe still has) a max number of dimensions.

Gaming

Programming

Assigning slice of a tensor to a variable?