def eye(N, M, k): return np.eye( numba_basic.to_scalar(N), numba_basic.to_scalar(M), numba_basic.to_scalar(k), dtype=dtype, )
def arange(start, stop, step): return np.arange( numba_basic.to_scalar(start), numba_basic.to_scalar(stop), numba_basic.to_scalar(step), dtype=dtype, )
def clip(_x, _min, _max): x = numba_basic.to_scalar(_x) _min_scalar = numba_basic.to_scalar(_min) _max_scalar = numba_basic.to_scalar(_max) if x < _min_scalar: return _min_scalar elif x > _max_scalar: return _max_scalar else: return x
def filldiagonaloffset(a, val, offset): height, width = a.shape if offset >= 0: start = numba_basic.to_scalar(offset) num_of_step = min(min(width, height), width - offset) else: start = -numba_basic.to_scalar(offset) * a.shape[1] num_of_step = min(min(width, height), height + offset) step = a.shape[1] + 1 end = start + step * num_of_step b = a.ravel() b[start:end:step] = val # TODO: This isn't implemented in Numba # a.flat[start:end:step] = val # return a return b.reshape(a.shape)
def broadcast_to(x, *shape): scalars_shape = create_zeros_tuple() i = 0 for s_i in literal_unroll(shape): scalars_shape = numba_basic.tuple_setitem( scalars_shape, i, numba_basic.to_scalar(s_i) ) i += 1 return np.broadcast_to(x, scalars_shape)
def careduce_axis(x): res_shape = res_shape_tuple_ctor(x.shape) x_axis_first = x.transpose(reaxis_first) res = np.full(res_shape, numba_basic.to_scalar(identity), dtype=dtype) for m in range(x.shape[axis]): reduce_fn(res, x_axis_first[m], res) return set_out_dims(res)
def join(axis, *tensors): return np.concatenate(tensors, numba_basic.to_scalar(axis))
def bartlett(x): return np.bartlett(numba_basic.to_scalar(x))
def careduce_axis(x): res = numba_basic.to_scalar(identity) for val in x: res = reduce_fn(res, val) return set_out_dims(res)