How do I replicate torch.chunk?

[vedantroy]

I have the following 2 lines of torch:

emb = emb[..., None, None]
cond_w, cond_b = th.chunk(emb, 2, dim=1)

I'm trying to convert this to einops.
So far I have:

emb = rearrange(emb, "b (c h w) -> b c h w", w=1, h=1)
cond_w, cond_b = th.chunk(emb, 2, dim=1)

which works pretty well.
But I can't figure out how to convert the last operation to einops.

I tried:

cond_w, cond_b = rearrange(emb, "b (split c) ... -> b split c ...", split=2)

but that is incorrect.

Here's my testing code:

def test_chunk():
    test_chunk_impl(th.Tensor([[1, 2, 3, 4], [4, 5, 6, 7]]).int())

def test_chunk_impl(a):
    b_t = a[..., None, None]
    b_e = rearrange(a, "b (c h w) -> b c h w", w=1, h=1)
    assert (b_t == b_e).all()

    c_t1, c_t2 = th.chunk(b_t, 2, dim=1)
    c_e1, c_e2 = rearrange(b_t, "b (split c) ... -> b split c ...", split=2)
    print(c_t1)
    print(c_e1)

    assert (c_t1 == c_e1).all()
    assert (c_t2 == c_e2).all()

test_chunk()

[vedantroy]

Solution:

    c_e1, c_e2 = rearrange(b_t, "b (split c) h w -> split b c h w", split=2)