Using vectorize on a function that expands the number of dimensions of the argument (or that returns multiple values)

0 views
Skip to first unread message

E. Tadeu

unread,
Jul 19, 2016, 8:54:58 AM7/19/16
to Numba Public Discussion - Public

Hi,

I’m evaluating the usage of Numba for production code, it’s impressive to see how it is evolving quickly.

I have a newbie question: is there a way to use @vectorize on a function that expands the number of dimensions of its arguments? For example, if a scalar is passed, it will return an 1D array; if an 1D array is passed, it will return a 2D array, etc.

Here is a very simplified example of what I’m trying to accomplish:

import numpy as np
import numba as nb

def f_pure(x):
    r = np.full((3,) + np.shape(x), np.nan)
    r[0, ...] = x
    r[1, ...] = x * 2.0
    return r

@nb.vectorize('float64(float64)')
def f_numba(x):
    r = np.full(3, np.nan)
    r[0] = x
    r[1] = x * 2.0
    return r

# NotImplementedError: array(float64, 1d, A) cannot be represented as a Numpy dtype
# @nb.vectorize('float64[:](float64)')
# def f_numba_2(x):
#     r = np.full(3, np.nan)
#     r[0] = x
#     r[1] = x * 2.0
#     return r

if __name__ == '__main__':
    a0 = 10.0
    a1 = np.array([1.0, 2.0])
    a2 = np.array([[0.5, 0.0, 2.0], [-1.0, -10.0, -100.0]])
    print(f_pure(a0))
    print(f_pure(a1))
    print(f_pure(a2))

    # TypeError: bad argument type for built-in operation
    # print(f_numba(10.0))

    # NotImplementedError: array(float64, 1d, A) cannot be represented as a Numpy dtype
    # print(f_numba_2(10.0))

I’m trying to replicate what f_pure is doing, but using @vectorize on a more simplified code (it can also be thought as a code taking scalars and returning multiple values).

Thanks!


Stephan Hoyer

unread,
Jul 19, 2016, 12:20:15 PM7/19/16
to numba...@continuum.io
I think this can be done with @guvectorize, if you write a wrapper that preallocates the result and writes to one of the inputs. (Otherwise, Numba can't figure out the correct size for the output array.)

--
You received this message because you are subscribed to the Google Groups "Numba Public Discussion - Public" group.
To unsubscribe from this group and stop receiving emails from it, send an email to numba-users...@continuum.io.
To post to this group, send email to numba...@continuum.io.
To view this discussion on the web visit https://groups.google.com/a/continuum.io/d/msgid/numba-users/CABTZkpuYAdUrRe%2Bvu_OpMmhVQ%2BqfjyDcGXBou77C1AAd1UU47Q%40mail.gmail.com.
For more options, visit https://groups.google.com/a/continuum.io/d/optout.

Reply all
Reply to author
Forward
0 new messages