def test_check_args_collapse_1():
ac, ag = gen_gpuarray((50, 1, 20), 'float32', ctx=context)
bc, bg = gen_gpuarray((50, 1, 20), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, bg), collapse=False)
assert n == 1000
assert nd == 3
assert dims == (50, 1, 20)
assert strs == ((80, 80, 4), (80, 80, 4))
assert offsets == (0, 0)
n, nd, dims, strs, offsets = check_args((ag, bg), collapse=True)
assert n == 1000
assert nd == 1
assert dims == (1000, )
assert strs == ((4, ), (4, ))
assert offsets == (0, 0)
def test_check_args_collapse_4():
ac, ag = gen_gpuarray((1, ), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, ), collapse=False)
assert n == 1
assert nd == 1
assert dims == (1, )
assert strs == ((4, ), )
assert offsets == (0, )
ac, ag = gen_gpuarray((1, 1), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, ), collapse=True)
assert n == 1
assert nd == 1
assert dims == (1, )
assert strs == ((4, ), )
assert offsets == (0, )
def test_check_args_collapse_1():
ac, ag = gen_gpuarray((50, 1, 20), 'float32', ctx=context)
bc, bg = gen_gpuarray((50, 1, 20), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, bg), collapse=False)
assert n == 1000
assert nd == 3
assert dims == (50, 1, 20)
assert strs == ((80, 80, 4), (80, 80, 4))
assert offsets == (0, 0)
n, nd, dims, strs, offsets = check_args((ag, bg), collapse=True)
assert n == 1000
assert nd == 1
assert dims == (1000,)
assert strs == ((4,), (4,))
assert offsets == (0, 0)
def test_check_args_collapse_4():
ac, ag = gen_gpuarray((1,), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag,), collapse=False)
assert n == 1
assert nd == 1
assert dims == (1,)
assert strs == ((4,),)
assert offsets == (0,)
ac, ag = gen_gpuarray((1, 1), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag,), collapse=True)
assert n == 1
assert nd == 1
assert dims == (1,)
assert strs == ((4,),)
assert offsets == (0,)
def test_check_args_broadcast_1():
ac, ag = gen_gpuarray((1, ), 'float32', ctx=context)
bc, bg = gen_gpuarray((50, ), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, bg), broadcast=True)
assert n == 50
assert nd == 1
assert dims == (50, )
assert strs == ((0, ), (4, ))
assert offsets == (0, 0)
def test_check_args_simple():
ac, ag = gen_gpuarray((50, ), 'float32', ctx=context)
bc, bg = gen_gpuarray((50, ), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, bg))
assert n == 50
assert nd == 1
assert dims == (50, )
assert strs == ((4, ), (4, ))
assert offsets == (0, 0)
ac, ag = gen_gpuarray((50, 1, 20), 'float32', ctx=context)
bc, bg = gen_gpuarray((50, 1, 20), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, bg))
assert n == 1000
assert nd == 3
assert dims == (50, 1, 20)
assert strs == ((80, 80, 4), (80, 80, 4))
assert offsets == (0, 0)
def test_check_args_broadcast_3():
ac, ag = gen_gpuarray((10, 20, 30), 'float32', ctx=context)
bc, bg = gen_gpuarray((1, 1, 1), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, bg), broadcast=True)
assert n == 6000
assert nd == 3
assert dims == (10, 20, 30)
assert strs == ((2400, 120, 4), (0, 0, 0))
assert offsets == (0, 0)
def test_check_args_simple():
ac, ag = gen_gpuarray((50,), 'float32', ctx=context)
bc, bg = gen_gpuarray((50,), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, bg))
assert n == 50
assert nd == 1
assert dims == (50,)
assert strs == ((4,), (4,))
assert offsets == (0, 0)
ac, ag = gen_gpuarray((50, 1, 20), 'float32', ctx=context)
bc, bg = gen_gpuarray((50, 1, 20), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, bg))
assert n == 1000
assert nd == 3
assert dims == (50, 1, 20)
assert strs == ((80, 80, 4), (80, 80, 4))
assert offsets == (0, 0)
def test_check_args_broadcast_3():
ac, ag = gen_gpuarray((10, 20, 30), 'float32', ctx=context)
bc, bg = gen_gpuarray((1, 1, 1), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, bg), broadcast=True)
assert n == 6000
assert nd == 3
assert dims == (10, 20, 30)
assert strs == ((2400, 120, 4), (0, 0, 0))
assert offsets == (0, 0)
def test_check_args_broadcast_1():
ac, ag = gen_gpuarray((1,), 'float32', ctx=context)
bc, bg = gen_gpuarray((50,), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, bg), broadcast=True)
assert n == 50
assert nd == 1
assert dims == (50,)
assert strs == ((0,), (4,))
assert offsets == (0, 0)
def test_check_args_collapse_3():
ac, ag = gen_gpuarray((50, 2, 10), 'float32', ctx=context, sliced=2,
offseted_outer=True)
bc, bg = gen_gpuarray((50, 2, 10), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, bg), collapse=True)
assert n == 1000
assert nd == 2
assert dims == (50, 20)
assert strs == ((160, 4), (80, 4))
assert offsets == (80, 0)
def test_check_args_broadcast_2():
ac, ag = gen_gpuarray((50, 1, 20), 'float32', ctx=context, sliced=2,
offseted_inner=True)
bc, bg = gen_gpuarray((50, 1, 20), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, bg), collapse=True,
broadcast=True)
assert n == 1000
assert nd == 2
assert dims == (50, 20)
assert strs == ((168, 4), (80, 4))
assert offsets == (4, 0)
def test_check_args_collapse_3():
ac, ag = gen_gpuarray((50, 2, 10),
'float32',
ctx=context,
sliced=2,
offseted_outer=True)
bc, bg = gen_gpuarray((50, 2, 10), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, bg), collapse=True)
assert n == 1000
assert nd == 2
assert dims == (50, 20)
assert strs == ((160, 4), (80, 4))
assert offsets == (80, 0)
def test_check_args_broadcast_2():
ac, ag = gen_gpuarray((50, 1, 20),
'float32',
ctx=context,
sliced=2,
offseted_inner=True)
bc, bg = gen_gpuarray((50, 1, 20), 'float32', ctx=context)
n, nd, dims, strs, offsets = check_args((ag, bg),
collapse=True,
broadcast=True)
assert n == 1000
assert nd == 2
assert dims == (50, 20)
assert strs == ((168, 4), (80, 4))
assert offsets == (4, 0)
def elemwise_collapse(dtype1, dtype2, shape1, shape2, expected):
assert len(shape1) == len(shape2)
# int8 does not cause problematic upcasts
scalar = numpy.asarray(1, dtype='int8')
a_cpu, a_gpu = gen_gpuarray(shape1, dtype1, ctx=context)
b_cpu, b_gpu = gen_gpuarray(shape2, dtype2, ctx=context)
o_shape = []
for i in range(len(shape1)):
o_shape.append(max(shape1[i], shape2[i]))
o = gpuarray.empty(o_shape, dtype=(a_cpu + b_cpu).dtype, context=context)
n, nd, dims, strs, offsets, contig = check_args((a_gpu, b_gpu),
collapse=True,
broadcast=True)
assert nd == expected, (shape1, shape2, dims, nd, expected)
k = ElemwiseKernel(context, [ArrayArg(numpy.dtype(dtype1), 'a'),
ArrayArg(numpy.dtype(dtype2), 'b'),
ArrayArg(o.dtype, 'o')], "o[i] = a[i] + b[i]")
out_cpu = a_cpu + b_cpu
k(a_gpu, b_gpu, o, collapse=True, broadcast=True)
assert numpy.allclose(numpy.asarray(o), out_cpu)
k(a_gpu, b_gpu, o, collapse=False, broadcast=True)
assert numpy.allclose(numpy.asarray(o), out_cpu)
broadcast = any([True for i in shape1 + shape2
if i == 1])
n, nd, dims, strs, offsets, contig = check_args((a_gpu, b_gpu, scalar),
collapse=True,
broadcast=True)
assert nd == expected
k = ElemwiseKernel(context, [ArrayArg(numpy.dtype(dtype1), 'a'),
ArrayArg(numpy.dtype(dtype2), 'b'),
ScalarArg(scalar.dtype, 's'),
ArrayArg(o.dtype, 'o')],
"o[i] = a[i] + b[i] + s")
out_cpu = a_cpu + b_cpu + scalar
k(a_gpu, b_gpu, scalar, o, collapse=True, broadcast=True)
assert numpy.allclose(numpy.asarray(o), out_cpu)
k(a_gpu, b_gpu, scalar, o, collapse=False, broadcast=True)
assert numpy.allclose(numpy.asarray(o), out_cpu)
if expected == 1:
expected2 = 2
else:
expected2 = expected
if len(shape1) != 4:
return
if shape1[0] != 1:
c_cpu, c_gpu = gen_gpuarray(shape1, dtype=dtype1, sliced=2, ctx=context)
n, nd, dims, strs, offsets,contig = check_args((c_gpu, b_gpu),
collapse=True,
broadcast=True)
if broadcast:
assert nd >= expected
else:
assert nd == expected2
