def test_multiple_shard_matrix(self): pwex = pywren.default_executor() X = np.random.randn(128, 128) X_sharded = local_numpy_init(X, (64, 64)) X_sharded_down = matrix_init.reshard_down(X_sharded, (4, 4), pwex) X_sharded_local = X_sharded.numpy() X_sharded_local_2 = X_sharded_down.numpy() X_sharded.free() X_sharded_down.free() assert (np.all(X_sharded_local == X)) assert (np.all(X_sharded_local_2 == X))
def test_single_shard_matrix(self): # pwex = pywren.default_executor() pwex = lithops.FunctionExecutor() X = np.random.randn(128,128) X_sharded = local_numpy_init(X, (128,128)) X_sharded_down = matrix_init.reshard_down(X_sharded, (4,4), pwex) X_sharded_local = X_sharded.numpy() X_sharded_local_2 = X_sharded_down.numpy() X_sharded.free() X_sharded_down.free() assert(np.all(X_sharded_local == X)) assert(np.all(X_sharded_local_2 == X))
def test_multiple_shard_tensor(self): pwex = pywren.default_executor() X = np.random.randn(128,128, 4) X_sharded = local_numpy_init(X, (64,64, 4)) X_sharded.autosqueeze = False X_sharded_down = matrix_init.reshard_down(X_sharded, (4,4,2), pwex=None) X_sharded_local = X_sharded.numpy() X_sharded_local_2 = X_sharded_down.numpy() X_sharded.free() X_sharded_down.free() X_sharded_down.get_block(0,0,0) assert(np.all(X_sharded_local == X)) assert(np.all(X_sharded_local_2 == X))