updated arange to mgrid in tutorials and references

[mm-aws06]

Issue #, if available:

<Fill ME>

Description of changes:

Updated arange to mgrid in tutorials and references

Testing:

Please see detailed unit test requirements in the CONTRIBUTING.md

• The change is covered by numeric check using nki.baremetal
• The change is covered by performance benchmark test using nki.benchmark
• The change is covered by end-to-end integration test

Pull Request Checklist

• I have filled in all the required field in the template
• I have tested locally that all the tests pass
• By submitting this pull request, I confirm that my contribution is made under the terms of the MIT-0 license.

[aws-qieqingy]

Please check the AutoTest error message, this line seems to have indentation error.

[mm-aws06]

fixed indentation and pushed again :)

[mm-aws06]

After reviewing my PR results, there's something I'd like to ask for anyone's opinion on. The changes I made from arange to mgrid in allocated_attention.py are correct (from what I can tell), however the indexing process is different:

arange and mgrid are both used to generate ranges of indices, but they return different shapes and types:
arange(N) returns a 1D array: [0, 1, ..., N-1]
mgrid[0:N, 0:M] returns two 2D arrays: one for row indices, one for column indices, both of shape (N, M)

Should I switch back from mgrid to arange here? Or alter the indexing? (and for the latter, will that involve making changes to the test files as well?)

[aws-qieqingy]

Hi @mm-aws06! The arange and mgrid has different meanings.

When indexing a tensor of shape (N, M) like the following, it can be viewed that we are using a compact loop representation.

t = nl.ndarray((N, M), ...)

ix = arange(N)[:, None] # shape is [N, 1]
iy = arange(M)[None, :] # shape is [1, M] 

t[ix, iy]

# is equivalent to

for i in range(N):
  for j in range(M):
    access t[ix[i, 0], j[0, j]] 

The [:, None] helps the compiler figure out which dimension is the index on.

For mgrid, it directly returns the index on both dimensions. It can still be viewed that we are using a compact loop representation, but the index has been already broadcasted for you by mgrid.

t = nl.ndarray((N, M), ...)

ix, iy = nl.mgrid[0:N, 0:M]

t[ix, iy]

for i in range(N):
  for j in range(M):
    access t[ix[i, j], iy[i, j]] 

Thanks for raising the issue. We will look into provide more documentation on the usage of mgrid and arange.