def METHOD_NAME(self, split_shape):
    custom_kernel = cupy.ElementwiseKernel(
        'float32 x, float32 y',
        'float32 z',
        'z = (x - y) * (x - y)',
        'custom')
    # TODO parameterize split shape
    np_a = numpy.arange(64).reshape(8, 8).astype(numpy.float32)
    np_b = (numpy.arange(64).reshape(8, 8) * 2.0).astype(numpy.float32)
    np_r = (np_a - np_b) * (np_a - np_b)
    d_a = cupyx.distributed._array.array(np_a, (0, 1), split_shape)
    d_b = cupyx.distributed._array.array(np_b, (0, 1), split_shape)
    d_r = custom_kernel(d_a, d_b)
    testing.assert_array_almost_equal(d_r.asnumpy(), np_r)