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

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

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

    assert agg.shape == (6, 16)
    assert agg_expr.isidentical(agg.sum(axis=0))

    for func in [var, std, mean]:
        (chunk, chunk_expr), (agg, agg_expr) = split(
            x, func(x, axis=0), chunk=c)

        assert chunk.shape in (4, 4)
        assert chunk_expr.shape == (1, 4)
        assert agg.shape == (6, 16)
dataset/ETHPy150Open blaze/blaze/blaze/expr/tests/test_split.py/test_nd_chunk_axis_args/operator `==` -> `in`