def test_nd_chunk():
    c = symbol('c', '4 * 4 * int32')

    (chunk, chunk_expr), (agg, agg_expr) = split(x, x.sum(), chunk=c)

    assert chunk.shape == (4, 4)
    assert chunk_expr.isidentical(chunk.sum(keepdims=True))

    assert agg.shape == (6, 4)
    assert agg_expr.isidentical(agg.sum())
dataset/ETHPy150Open blaze/blaze/blaze/expr/tests/test_split.py/test_nd_chunk/original