# Path: taskloaf/csr.py
# def distribute(mat, gang):
#     n_chunks = len(gang)
#     nrows = mat.shape[0]
#     ncols = mat.shape[1]
#     print(n_chunks, nrows, ncols)
# 
#     # TODO: optimal chunking?
#     row_chunk_bounds = np.linspace(0, nrows, n_chunks + 1).astype(np.int)
#     dist_chunks = []
#     for i in range(n_chunks):
#         start_row = row_chunk_bounds[i]
#         end_row = row_chunk_bounds[i + 1]
#         chunk_mat = mat[start_row:end_row]
#         dist_chunks.append(TskCSRChunk(chunk_mat, start_row))
# 
#     return TskCSR(mat.shape, dist_chunks, gang)
# 
# class TskArray:
#     def __init__(self, vals=None, size=None):
#         if vals is not None:
#             assert len(vals.shape) == 1
#             self.dtype = vals.dtype
#             self.ref = tsk.alloc(vals.shape[0] * self.dtype.itemsize)
#             self._array(self.ref.get_local())[:] = vals
#         else:
#             self.dtype = np.float64()
#             self.ref = tsk.alloc(size * self.dtype.itemsize)
# 
#     def _array(self, buf):
#         return np.frombuffer(buf, dtype=self.dtype)
# 
#     async def array(self):
#         return self._array(await self.ref.get())


import logging
import numpy as np
import scipy.sparse
import taskloaf as tsk
from taskloaf.csr import distribute, TskArray
    return A


def setup_worker(name, cfg):
    tsk.cfg.stdout_logging(name, logger_name="taskloaf.profile")


def main():
    cfg = tsk.Cfg(
        n_workers=2, log_level=logging.WARN, initializer=setup_worker
    )

    async def f():
        nrows = int(1e7)
        nnz = nrows * 10
        n_repeats = 1
        mat = random_test_matrix(nrows, nnz)
        vec = np.random.rand(nrows) - 0.5
        t = tsk.Timer()
        for i in range(n_repeats):
            correct = mat.dot(vec)
        t.report("simple dot")

        gang = await tsk.ctx().wait_for_workers(cfg.n_workers)
        t.report("wait for workers")

        t.report("launch profiler")
        tsk_vec = TskArray(vals=vec)
        t.report("shmem v")


# Path: taskloaf/csr.py
# def distribute(mat, gang):
#     n_chunks = len(gang)
#     nrows = mat.shape[0]
#     ncols = mat.shape[1]
#     print(n_chunks, nrows, ncols)
# 
#     # TODO: optimal chunking?
#     row_chunk_bounds = np.linspace(0, nrows, n_chunks + 1).astype(np.int)
#     dist_chunks = []
#     for i in range(n_chunks):
#         start_row = row_chunk_bounds[i]
#         end_row = row_chunk_bounds[i + 1]
#         chunk_mat = mat[start_row:end_row]
#         dist_chunks.append(TskCSRChunk(chunk_mat, start_row))
# 
#     return TskCSR(mat.shape, dist_chunks, gang)
# 
# class TskArray:
#     def __init__(self, vals=None, size=None):
#         if vals is not None:
#             assert len(vals.shape) == 1
#             self.dtype = vals.dtype
#             self.ref = tsk.alloc(vals.shape[0] * self.dtype.itemsize)
#             self._array(self.ref.get_local())[:] = vals
#         else:
#             self.dtype = np.float64()
#             self.ref = tsk.alloc(size * self.dtype.itemsize)
# 
#     def _array(self, buf):
#         return np.frombuffer(buf, dtype=self.dtype)
# 
#     async def array(self):
#         return self._array(await self.ref.get())


# Path: examples/new_sparse.py
import logging
import numpy as np
import scipy.sparse
import taskloaf as tsk
from taskloaf.csr import distribute, TskArray

    return A


def setup_worker(name, cfg):
    tsk.cfg.stdout_logging(name, logger_name="taskloaf.profile")


def main():
    cfg = tsk.Cfg(
        n_workers=2, log_level=logging.WARN, initializer=setup_worker
    )

    async def f():
        nrows = int(1e7)
        nnz = nrows * 10
        n_repeats = 1
        mat = random_test_matrix(nrows, nnz)
        vec = np.random.rand(nrows) - 0.5
        t = tsk.Timer()
        for i in range(n_repeats):
            correct = mat.dot(vec)
        t.report("simple dot")

        gang = await tsk.ctx().wait_for_workers(cfg.n_workers)
        t.report("wait for workers")

        t.report("launch profiler")
        tsk_vec = TskArray(vals=vec)
        t.report("shmem v")


        tsk_mat = distribute(mat, gang)