search

modularml / mojo

The Mojo Programming Language
https://docs.modular.com/mojo/manual/
Other
23.08k stars 2.59k forks source link

[Feature Request] Infer `NDBuffer` Rank From DimList #4

Open lsh opened 1 year ago

lsh commented 1 year ago

In most APIs, the rank of a multidimensional array would be assumed to be equivalent to the shape.

I noticed that the invocation for NDBuffer is NDBuffer[rank, shape, dtype](). It would be ergonomic to have a constructor that assumes rank from shape, so it could just be NDBuffer[DimList(2,3), DType.ui32](), or even better NDBuffer[(2,3), DType.ui32]()

Mogball commented 1 year ago

Thanks for the issue!

This is a current limitation of how we handle dependent parameters, and there are other types and functions that suffer from this. The key issue is that the shape parameter depends on rank, and parameters are resolved left-to-right, so rank has to be defined first. That said, I think this is something we should support, since it is natural and intuitive to infer rank from shape.

lattner commented 1 year ago

@abduld this seems like we can just directly improve NDBuffer. Could we change it from:

struct NDBuffer[
    rank: Int,
    shape: DimList,
    type: DType,
]:

to:

struct NDBuffer[
    shape: DimList,
    type: DType,
]:
  alias rank = len(shape) 
  ...

DimList doesn't currently have a __len__ method, but its underlying storage type does.

abduld commented 1 year ago

Yes, we can certainly improve (and will improve) the API for NDBuffer. Thanks @lsh for the suggestion