def kron(a, b):
# pylint: disable=protected-access,g-complex-comprehension
a, b = array_ops._promote_dtype(a, b)
ndim = max(a.ndim, b.ndim)
if a.ndim < ndim:
a = array_ops.reshape(a, array_ops._pad_left_to(ndim, a.shape))
if b.ndim < ndim:
b = array_ops.reshape(b, array_ops._pad_left_to(ndim, b.shape))
a_reshaped = array_ops.reshape(a, [i for d in a.shape for i in (d, 1)])
b_reshaped = array_ops.reshape(b, [i for d in b.shape for i in (1, d)])
out_shape = tuple(np.multiply(a.shape, b.shape))
return array_ops.reshape(a_reshaped * b_reshaped, out_shape)
def reshape(x, num_devices=2):
x_shape = list(x.shape)
batch_size = x_shape[0]
batch_size_per_device = batch_size // num_devices
# New shape.
new_shape_prefix = [num_devices, batch_size_per_device]
return array_ops.reshape(x, new_shape_prefix + x_shape[1:])
def run_test(arr, newshape, *args, **kwargs):
for fn1 in self.array_transforms:
for fn2 in self.array_transforms:
arr_arg = fn1(arr)
newshape_arg = fn2(newshape)
self.match(
array_ops.reshape(arr_arg, newshape_arg, *args,
**kwargs),
np.reshape(arr_arg, newshape, *args, **kwargs))