fix: consistent shape convention in flip/reshape/broadcast_to/stack

[henryiii]

🤖 AI text below 🤖

Part of #144.

The inconsistency

ragged.array.shape uses None for ragged (variable-length) dimensions and ints for regular dimensions. Several manipulation functions have a numpy "fast path" for uniform inputs and a general "ragged path" for everything else. For four of them, the two paths disagreed on the shape signature: the fast path made every inner dimension a fixed int, while the general path left the source array's own inner dimensions variable-length (None).

This meant the same logical operation returned a different shape signature depending only on whether the data happened to be uniform — a silent behavioral fork. (permute_dims and roll were already correct: they route their fast path through the existing _ak_from_numpy helper.)

For x = ragged.array([[1.0, 2.0], [3.0, 4.0]]) with x.shape == (2, None):

operation	before (uniform fast path)	after (this PR)	general/ragged path
flip(x, axis=1)	(2, 2)	(2, None)	(2, None)
flip(x)	(2, 2)	(2, None)	(2, None)
reshape(x, (2, 2))	(2, 2)	(2, None)	n/a (ragged reshape is 1-D only)
reshape(x, (2, 2, 1))	(2, 2, 1)	(2, None, None)	n/a
broadcast_to(x, (2, 2))	(2, 2)	(2, None)	(2, None)
stack([x, x])	(2, 2, 2)	(2, 2, None)	(2, 2, None)
stack([x, x], axis=2)	(2, 2, 2)	(2, None, 2)	(2, None, 2)

1-D inputs are unaffected (they have no inner ragged dimension): reshape(v, (3,)), flip(v), broadcast_to(v, (2, 3)) -> (2, None), stack([v, v]) -> (2, 3).

Approach per function

I verified the general path's output per function rather than assuming a single fix:

• flip — route the numpy fast path through _ak_from_numpy instead of plain ak.from_numpy. flip preserves the input shape, so a (2, None) input must stay (2, None).
• reshape — same: route the multi-dimensional fast path through _ak_from_numpy. The reshaped result of a ragged-typed array is itself a transformed ragged array, so inner dims become None. (The ragged path only ever produces 1-D output, so there is no ragged equivalent for multi-dim reshape to diff against; the convention matches the rest of the library.)
• broadcast_to — _ak_from_numpy alone is not correct here, because the general path keeps newly prepended outer dims regular while making the source's own inner dims None (e.g. (2, 3)-uniform → (4, 2, 3) target → (4, 2, None)). I confirmed the awkward broadcast_arrays general path already produces exactly this convention for every test case (values identical), so I removed the numpy fast path and let the general path run unconditionally.
• stack — I ran both paths and compared. The general path (expand_dims + concat) keeps the newly inserted stack axis regular while making the source dims None (axis 0 → (2, 2, None), axis 2 → (2, None, 2)). Plain _ak_from_numpy would wrongly make all inner dims None ((2, None, None)), so I removed the numpy fast path and let the expand_dims + concat general path run unconditionally. Added a parametrized test asserting the uniform and genuinely-ragged inputs produce the same shape structure for every axis.

Boundary: what is intentionally not changed

meshgrid and the creation functions (zeros, ones, etc.) are left unchanged. They create new regular data (e.g. ragged.zeros((2, 2)) legitimately has int dims) rather than transforming an existing ragged array, so the int-dim convention is correct for them.

Existing tests changed (user-visible behavior change)

These assertions pinned the old int-shape convention and were updated (with an explanatory comment in each):

• tests/test_broadcast_to.py::TestBroadcastTo2DUniform::test_outer_replication — (4, 2, 3) → (4, 2, None) (2-D source's inner dim is variable-length).
• tests/test_broadcast_to.py::TestBroadcastTo3DUniform::test_outer_replication — (5, 2, 3, 4) → (5, 2, None, None) (3-D source is (2, None, None)).
• tests/test_reshape.py::TestReshapeRegular::test_add_trailing_dim — (2, 2, 1) → (2, None, None).

The stack shape tests already asserted the (2, 2, None) convention (those were the ragged-path tests), so no stack assertions needed changing — the uniform fast path now simply agrees with them.

New tests pinning the convention

Added a …ShapeConvention class to each of test_flip.py, test_reshape.py, test_broadcast_to.py, and test_stack.py asserting that a uniform-but-ragged-typed input produces the same shape signature as the general/ragged path (including a parametrized stack test that diffs the uniform vs genuinely-ragged result for every axis).

Heads-up for maintainers

This is a user-visible behavior change to the returned shape of flip, reshape, broadcast_to, and stack on multi-dimensional inputs (values are unchanged; only the int-vs-None dimension signature changes). It is opened as a draft for discussion. One judgment call worth flagging: for broadcast_to with an explicit all-int target shape, one could argue the result should be regular; I chose to match the ragged path (inner source dims → None) for consistency, since the same operation on a genuinely ragged source can only produce None there anyway.

Full suite: 1313 passed, 6 skipped (pre-existing numpy-v1-only skips). prek -a (ruff + ruff-format + mypy --strict) clean.

🤖 Generated with Claude Code