How can i return from a function an array or slice with different dimension like [][]T or [][][]T

I have a struct called Tensor that has a shape attribute and a data attribute defined as data: []T that stores data in a linear way. I want to do a method called ToArray that returns the linear data of the tensor struct as a multidim array. the problem is that i don’t know the return type of this method since shaper could be 2,3,4 or more dimension and i can’t do toArray(self: *@This()) anytype{} how would you do this?

Tks a lot for the help I’m loving this language

You can use comptime to dynamically choose the appropriate return type at compile time, e.g.:

fn MagicalReturnType(comptime DataType: type, ...) type {
    // ...your magic here...
}

pub fn Tensor(comptime T: type) type {
    return struct {
        data: []T,

        pub fn toArray(self: @This(), ...) MagicalReturnType([]T, ...) {
            // ...
        }
    };
};
1 Like

I’ve tried with

pub fn toArray(self: @This()) MagicalReturnType(T, self.shape.len) {
            if (self.shape.len == 1) {
                return self.data;
            }
            return constructMultidimensionalArray(T, self.data, self.shape, 0);
        }

        // create multi dim array
        fn constructMultidimensionalArray(comptime ElementType: type, data: []ElementType, shape: []usize, depth: usize) MagicalReturnType(ElementType, shape.len - depth) {
            if (depth == shape.len - 1) {
                return data;
            }

            const current_dim = shape[depth];
            var result = try std.heap.page_allocator.alloc(MagicalReturnType(ElementType, shape.len - depth - 1), current_dim);

            var offset: usize = 0;
            const sub_array_size = calculateProduct(shape[depth + 1 ..]);

            for (current_dim) |i| {
                result[i] = constructMultidimensionalArray(ElementType, data[offset .. offset + sub_array_size], shape, depth + 1);
                offset += sub_array_size;
            }

            return result;
        }


        fn calculateProduct(slice: []usize) usize {
            var product: usize = 1;
            for (slice) |elem| {
                product *= elem;
            }
            return product;
        }


        fn MagicalReturnType(comptime DataType: type, comptime dim_count: usize) type {
            return if (dim_count == 1) []DataType else []MagicalReturnType(DataType, dim_count - 1);
        }

but it gives


`TheBigBook/tensor.zig:69:64: error: unable to evaluate comptime expression
        pub fn toArray(self: @This()) MagicalReturnType(T, self.shape.len) {
                                                           ~~~~^~~~~~`

comptime vs runtime

That is the main issue I see here, you need to decide between either handling the shape at comptime which would allow you to use different types depending on the shape picked at comptime, or handle it all at run time which means you need to pick a type that can handle all shapes.

I think one way you could do this is by renaming/repurposing toArray to toShape and adding a comptime shape parameter to it, which then is used to calculate the return type, within the body of the function you check at run time that the runtime shape matches the comptime shape, if it doesn’t you either error or panic.

Another problem is the whole issue of multidimensional slices (slices need their own storage and point to other storage combined with creating new data, you end up needing dynamic memory allocation), if you are going this route where the shape has to be provided as a comptime parameter, then there isn’t a good reason to use slices at all, because if the shape is known at comptime you can instead use multidimmensional arrays like in this answer:

If you don’t want to provide the shape as a comptime parameter, you could instead return a tagged union that enumerates the possible variants (effectively just delaying the choice by wrapping it in a type that can represent every choice), but if you want to support a wide combination of possible shapes that might end up too messy and difficult. Also in the end the union would mostly serve as a way to indicate at run time which of the shapes was chosen and then the person receiving the value would have to switch on the tagged union, selecting one of its variants (and at that point that code branch would again know the shape at comptime (because that branch only gets chosen for that variant)).

Overall I think it would be helpful if you described your overall goals and constraints, so that we can make more specific suggestions how you could reach your goal.

For example whether you really want to process all kinds of generic shapes, or more specific ones, whether you know which shapes will ever be needed at comptime, or your program only finds that out at run time.

How you need to access the data, what operations you need on the data, etc.

It is easier to deal with (and write efficient code for) specific comptime known shapes then having to deal with all possible shapes at run time.


You also might be interested in these projects:

I don’t understand the details of those projects well enough, but from what I have gathered they also deal with lots of different dimmensional data, so there may be something there that is related to what you are trying to do.

2 Likes

This is my code for a Tensor struct in Zig, where the data is stored in a 1D linear array ([]T). I would like to create a function toArray that uses the shape to return the data as a multidimensional array. The issue I am facing is that the return type could vary depending on the shape. For example, it could return a [][][]T, [][]T, or even higher dimensions, and I am unsure how to handle this dynamic return type because of the multiple possible array types.

pub fn Tensor(comptime T: type) type {
    return struct {
        data: []T,
        size: usize,
        shape: []usize,
        allocator: *const std.mem.Allocator,

        pub fn fromArray(allocator: *const std.mem.Allocator, inputArray: anytype, shape: []usize) !@This() {
            //std.debug.print("\n fromArray initialization...", .{});
            var total_size: usize = 1;
            for (shape) |dim| {
                total_size *= dim;
            }
            const tensorShape = try allocator.alloc(usize, shape.len);
            @memcpy(tensorShape, shape);

            const tensorData = try allocator.alloc(T, total_size);
            _ = flattenArray(T, inputArray, tensorData, 0);

            return @This(){
                .data = tensorData,
                .size = total_size,
                .shape = tensorShape,
                .allocator = allocator,
            };
        }

        pub fn init(allocator: *const std.mem.Allocator) !@This() {
            return @This(){
                .data = &[_]T{},
                .size = 0,
                .shape = &[_]usize{},
                .allocator = allocator,
            };
        }

        //copy self and return it in another Tensor
        pub fn copy(self: *@This()) !Tensor(T) {
            return try Tensor(T).fromArray(self.allocator, self.data, self.shape);
        }

        //inizialize and return a all-zero tensor starting from the shape
        pub fn fromShape(allocator: *const std.mem.Allocator, shape: []usize) !@This() {
            var total_size: usize = 1;
            for (shape) |dim| {
                total_size *= dim;
            }

            const tensorData = try allocator.alloc(T, total_size);
            for (tensorData) |*data| {
                data.* = 0;
            }

            return @This().fromArray(allocator, tensorData, shape);
        }

        //pay attention, the fill() can also perform a reshape
        pub fn fill(self: *@This(), inputArray: anytype, shape: []usize) !void {

            //deinitialize data e shape
            self.deinit(); //if the Tensor has been just init() this function does nothing

            //than, filling with the new values
            var total_size: usize = 1;
            for (shape) |dim| {
                total_size *= dim;
            }
            const tensorShape = try self.allocator.alloc(usize, shape.len);
            @memcpy(tensorShape, shape);

            const tensorData = try self.allocator.alloc(T, total_size);
            _ = flattenArray(T, inputArray, tensorData, 0);

            self.data = tensorData;
            self.size = total_size;
            self.shape = tensorShape;
        }

        pub fn deinit(self: *@This()) void {
            //std.debug.print("\n deinit tensor:\n", .{});
            // Verifica se `data` è valido e non vuoto prima di liberarlo
            if (self.data.len > 0) {
                //std.debug.print("Liberazione di data con lunghezza: {}\n", .{self.data.len});
                self.allocator.free(self.data);
                self.data = &[_]T{}; // Resetta lo slice a vuoto
            }
            // Verifica se `shape` è valido e non vuoto prima di liberarlo
            if (self.shape.len > 0) {
                //std.debug.print("Liberazione di shape con lunghezza: {}\n", .{self.shape.len});
                self.allocator.free(self.shape);
                self.shape = &[_]usize{}; // Resetta lo slice a vuoto
            }
        }

        pub fn setShape(self: *@This(), shape: []usize) !void {
            var total_size: usize = 1;
            for (shape) |dim| {
                total_size *= dim;
            }
            self.shape = shape;
            self.size = total_size;
        }

        pub fn getSize(self: *@This()) usize {
            return self.size;
        }

        pub fn get(self: *const @This(), idx: usize) !T {
            if (idx >= self.data.len) {
                return error.IndexOutOfBounds;
            }
            return self.data[idx];
        }

        pub fn set(self: *@This(), idx: usize, value: T) !void {
            if (idx >= self.data.len) {
                return error.IndexOutOfBounds;
            }
            self.data[idx] = value;
        }

        pub fn flatten_index(self: *const @This(), indices: []const usize) !usize {
            var idx: usize = 0;
            var stride: usize = 1;
            for (0..self.shape.len) |i| {
                idx += indices[self.shape.len - 1 - i] * stride;
                stride *= self.shape[self.shape.len - 1 - i];
            }
            return idx;
        }

        pub fn get_at(self: *const @This(), indices: []const usize) !T {
            const idx = try self.flatten_index(indices);
            return self.get(idx);
        }

        pub fn set_at(self: *@This(), indices: []const usize, value: T) !void {
            const idx = try self.flatten_index(indices);
            return self.set(idx, value);
        }

        pub fn info(self: *@This()) void {
            std.debug.print("\ntensor infos: ", .{});
            std.debug.print("\n  data type:{}", .{@TypeOf(self.data[0])});
            std.debug.print("\n  size:{}", .{self.size});
            std.debug.print("\n shape.len:{} shape: [ ", .{self.shape.len});
            for (0..self.shape.len) |i| {
                std.debug.print("{} ", .{self.shape[i]});
            }
            std.debug.print("] ", .{});
            self.print();
        }

        pub fn print(self: *@This()) void {
            std.debug.print("\n  tensor data: ", .{});
            for (0..self.size) |i| {
                std.debug.print("{} ", .{self.data[i]});
            }
            std.debug.print("\n", .{});
        }
    };
}

I would appreciate any advice on how to handle the return type dynamically based on the shape. Thanks for your help!

For efficiency reasons I would really suggest not to turn this into a multi-dimensional slice, because it requires a ton of allocations and each access has to go through multiple indirections to access the value. Writing tensor.get(.{x, y, z}) and tensor.set(.{x, y, z}, val) is really not that innconvenient as an alternative.

Ok now to actually solve your problem:
The problem here is that the dimension is not known at comptime.

But when you call toArray, you, the programmer, already know the dimension. Otherwise you wouldn’t be able to use of the resulting array (you can’t write array[x][y] without knowing that it’s 2 dimensions). Therefor, you could pass this dimension into the function as a comptime parameter and replace all usages of shapes.len with this parameter.
Additionally you need to make the depth comptime as well.

        pub fn toArray(self: @This(), comptime dimension: usize) MagicalReturnType(T, dimension) {
            std.debug.assert(self.shape.len == dimension);
            if (self.shape.len == 1) {
                return self.data;
            }
            return constructMultidimensionalArray(T, self.data, self.shape, 0, dimension);
        }

        fn constructMultidimensionalArray(comptime ElementType: type, data: []ElementType, shape: []usize, comptime depth: usize, comptime dimension: usize) MagicalReturnType(ElementType, dimension - depth) {
            if (depth == dimension - 1) {
                return data;
            }

            const current_dim = shape[depth];
            var result = try std.heap.page_allocator.alloc(MagicalReturnType(ElementType, dimension - depth - 1), current_dim);
            var offset: usize = 0;
            const sub_array_size = calculateProduct(shape[depth + 1 ..]);

            for (current_dim) |i| {
                result[i] = constructMultidimensionalArray(ElementType, data[offset .. offset + sub_array_size], shape, depth + 1, dimension);
                offset += sub_array_size;
            }

            return result;
        }
3 Likes

I just updated the code with all of my tensor class (look at the snippet), is it better in term of efficiency? tks a lot

1 Like

Yeah, it looks better.

1 Like