@pearu and I had a chat today about how we should support the different ways of manipulating xnd objects in numba.

We agreed that the API supported inside a jitted function should be the same as without the jit decorator. So we need to understand when operations in xnd return new xnd objects and when they return python objects. If a user calls an operation that returns a Python object inside a Numba function, we have to figure out the appropriate numba type for this, so that numba can lower the result.

Here are some of different examples of indexing and iterating, followed by some thoughts.

Indexing

First let's give some examples of indexing, which we could support in Numba. A user could iterate through a range of numbers, and index XND objects using those to get other xnd object or python objects

Arrays

By ints

...results in an array with one less dimension.

fixed arrays:

In [2]: x = xnd([[1, 2], [3, 4]]); x
Out[2]: xnd([[1, 2], [3, 4]], type='2 * 2 * int64')

In [3]: x[0]
Out[3]: xnd([1, 2], type='2 * int64')

Variable dimension arrays:

In [4]: x = xnd([[1], [1, 2]]); x
Out[4]: xnd([[1], [1, 2]], type='var * var * int64')

In [5]: x[0]
Out[5]: xnd([1], type='var * int64')

Except if the result would be a scalar, it is not returned as an xnd array, but just a Python scalar:

In [6]: x = xnd(1); x
Out[6]: xnd(1, type='int64')

In [8]: x[tuple()]
Out[8]: 1

In [9]: x = xnd([1, 2]); x
Out[9]: xnd([1, 2], type='2 * int64')

In [10]: x[0]
Out[10]: 1

By Tuples of ints

is equivalent to splitting out the indexing, i.e. x[i, j...] == x[i][j...]:

In [27]: x = xnd([[[1], [2]], [[3], [4]]]); x
Out[27]: xnd([[[1], [2]], [[3], [4]]], type='2 * 2 * 1 * int64')

In [28]: x[0, 1]
Out[28]: xnd([2], type='1 * int64')

In [29]: x[0, 1]
Out[29]: xnd([2], type='1 * int64')

By Slices

Results in same dimension array, with shape shrunk down:

In [58]: x = xnd([[1, 2], [3, 4]]); x
Out[58]: xnd([[1, 2], [3, 4]], type='2 * 2 * int64')

In [59]: x[1:]
Out[59]: xnd([[3, 4]], type='1 * 2 * int64')

In [60]: x[:-1]
Out[60]: xnd([[1, 2]], type='1 * 2 * int64')

In [61]: x[:]
Out[61]: xnd([[1, 2], [3, 4]], type='2 * 2 * int64')

In [62]: x[1:2]
Out[62]: xnd([[3, 4]], type='1 * 2 * int64')

With step, you can reverse/skip items:

In [74]: x = xnd([[1, 2], [3, 4], [5, 6]]); x
Out[74]: xnd([[1, 2], [3, 4], [5, 6]], type='3 * 2 * int64')

In [75]: x[::2]
Out[75]: xnd([[1, 2], [5, 6]], type='2 * 2 * int64')

In [76]: x[::-1]
Out[76]: xnd([[5, 6], [3, 4], [1, 2]], type='3 * 2 * int64')

Both of these work on variable dimension arrays:

In [90]: x = xnd([[1, 2], [3, 4, 5]]); x
Out[90]: xnd([[1, 2], [3, 4, 5]], type='var * var * int64')

In [91]: x[1:]
Out[91]: xnd([[3, 4, 5]], type='var * var * int64')

In [92]: x[::-1]
Out[92]: xnd([[3, 4, 5], [1, 2]], type='var * var * int64')

Tuples of slices

If you index by a slice and an integer, in a tuple, then each dimension that is indexed by a integer is removed, whereas each dimension that is sliced is just compressed:

In [95]: x = xnd([[1, 2], [3, 4]]); x
Out[95]: xnd([[1, 2], [3, 4]], type='2 * 2 * int64')

In [96]: x[::-1, 0]
Out[96]: xnd([3, 1], type='2 * int64')

In [97]: x[0, ::-1]
Out[97]: xnd([2, 1], type='2 * int64')

You can't do this mixed indexing with variable dimensions however:

In [90]: x = xnd([[1, 2], [3, 4, 5]]); x
Out[90]: xnd([[1, 2], [3, 4, 5]], type='var * var * int64')

In [93]: x[::-1, 0]
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-93-9b89208d65ed> in <module>()
----> 1 x[::-1, 0]

IndexError: mixed indexing and slicing is not supported for var dimensions

Tuples

by ints gets the ith element from the tuple:

In [102]: x = xnd((1, 2)); x
Out[102]: xnd((1, 2), type='(int64, int64)')

In [103]: x[0]
Out[103]: 1

Records

by strings gets the field from the struct:

In [114]: x = xnd({"hi": 1}); x
Out[114]: xnd({'hi': 1}, type='{hi : int64}')

In [116]: x['hi']
Out[116]: 1

strings/bytes

Are not indexable:

In [38]: x = xnd("hi"); x.type
Out[38]: ndt("string")

In [39]: x = xnd("hi"); x
Out[39]: xnd('hi', type='string')

In [40]: x[0]
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-40-2f755f117ac9> in <module>()
----> 1 x[0]

TypeError: type not indexable

In [35]: x = xnd("hi", type='fixed_string(2)'); x
Out[35]: xnd('hi', type='fixed_string(2)')

In [36]: x[0]
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-36-2f755f117ac9> in <module>()
----> 1 x[0]

TypeError: type not indexable

Categorical

Indexing an array of categoricals will give you back a python string:

In [119]: levels = ['January', 'August', 'December', None]

In [120]: x = xnd(['January', 'January', None, 'December', 'August', 'December', 'December'], levels=levels)

In [121]: x = xnd(['January', 'January', None, 'December', 'August', 'December', 'December'], levels=levels)

In [122]: x
Out[122]:
xnd(['January', 'January', None, 'December', 'August', 'December', 'December'],
    type='7 * categorical('January', 'August', 'December', NA)')

In [123]: x[0]
Out[123]: 'January'

Iterator

A user might also want to loop directly over xnd structures in Numba, without indexing.

Iterating through an xnd object is the same (in all cases I believe) to yielding a each item from 0 to the length of the object. i.e. iter(x) == (x[i] for i in range(len(x)):

In [130]: list(xnd([1, 2, 3]))
Out[130]: [1, 2, 3]

In [131]: list(xnd([[1], [2], [3]]))
Out[131]:
[xnd([1], type='1 * int64'),
 xnd([2], type='1 * int64'),
 xnd([3], type='1 * int64')]

In [132]: list(xnd((1, 2)))
Out[132]: [1, 2]

In [136]: list(xnd({"hi": 123, "there": 124}))
Out[136]: [123, 124]

Thoughts

I notice that xnd actually doesn't give you a 0D array when indexing 1D arrays, instead it gives you a python object. If the user never gets a scalar at the numba level, and instead just calls gumath kernels on higher dimensional types, then we don't need to know the type of each xnd object at compile time, because the gumath apply function takes care of understanding those things at runtime. This could look like:

@jit
def sum(a):
   return gumath.sum(a)

sum(xnd([1, 2, 3]))

However, if we do want to let users get access to the underlying python objects, then we need to have type information at compile time, so that numba knows what numba type to associate with the results of indexing. For example:

@jit(nopython=True)
def sum(a):
    s = 0
    for i in range(len(a)):
        s += a[i]
    return s

sum(xnd([1, 2, 3]))

Or matrix multiplication:

@jit(nopython=True)
def matrix_multiply(a, b, c):
    n, m = a.type.shape
    m, p = b.type.shape
    for i in range(n):
        for j in range(p):
            c[i][j] = 0
            for k in range(m):
                c[i][j] += a[i][k] * b[k][j]

matrix_multiply(xnd([[1, 2], [3, 4]]), xnd([[1, 2], [3, 4]]), xnd.empty("2 * 2 * int64"))

If we then can register this a gumath kernel then the gumath_apply can take care of creating the empty resulting shape:

@register_gumath_kernel('N * M * int64, M * P * int64 -> M * P * int64')
@jit(nopython=True)
def matrix_multiply(a, b, c):
    n, m = a.type.shape
    m, p = b.type.shape
    for i in range(n):
        for j in range(p):
            c[i][j] = 0
            for k in range(m):
                c[i][j] += a[i][k] * b[k][j]

matrix_multiply(xnd([[1, 2], [3, 4]]), xnd([[1, 2], [3, 4]]))

It's less clear how you would implement matrix multiple in Numba, with just gumath functions. But maybe this is a silly example, because you wouldn't want to implement matrix multiplication in Numba, you would want to use an existing implementation that is an existing gumath kernel.

Let's come back to why we are interesting in using Numba in the first place with xnd. My understanding is that it gives us a fast way of mixing gumath kernel application with control flow. Right now, if you want to add three gumath arrays, you would call gumath.add(gumath.add(a, b), c) and between those applications it has to come back to the python runtime and create a new xnd python object. If we did that in a numba jitted function, it would be able to stay out of python the whole time.

Same if we add some control flow to computation on xnd objects. For example, consider this silly example:

c = xnd(0)
for j in xnd[[1, 2], [3, 4]]):
    if j[0] > 1:
        c = gumath.add(c, gumath.sum(j))

Is this the sort of use case we are trying to optimize with Numba? Here, we have control flow dependent on the values inside xnd arrays, so we would have to expose them to Numba. Whereas, in the previous example of gumath.add(gumath.add(a, b), c), numba could be entirely oblivious to what is the ndtype of the xnd containers, because it just forwards them to gumath_apply.

I wonder how much we could get if we kept everything as just gumath kernel application. What if slicing and indexing where actually just gumath kernels, instead of being at the xnd level? So that they always returned xnd objects, never python objects. Could we type slicing in gumath?

If we could then pass gumath kernels themselves into gumath kernels, then could write all the control flow in kernels, like have an if kernel that takes takes a boolean scalar array and two other arrays. But then we have to have some sort of lazy evaluation/thunks, since if should only run one branch. Well then you could have a gumath kernel called thunk that takes in a gumath kernel, and some arguments, and returns another gumath kernel that takes no arguments and returns the result of calling that gumath kernel.

This seems like a complicated road to go down, to keep as much logic as possible in gumath kernels.

What if we take the opposite approach, and instead of making numba as unaware as possible, make it more aware about xnd types? We can store the xnd type on the numba type, so that at compile time we know the type of everything. With some indexing, we can still know the resulting type information at compile time (i.e. without knowing what the index is). For example 3 * 4 * int64 indexed by an int returns a shape 4 * int64. But, if we have {hi : int64, there: float32} and we index it, we don't know the type of the resulting value.

Numba must do something like this for Numpy however, so we should investigate how it works in that way. Does it always know the resulting numpy dtype and shape at compile time? Or are there times it does not? Or maybe there certain contexts for indexing where it can be lowered and other times it cannot?

xnd-project / numba-xnd

Different ways of indexing #5