Problem
The Python type stub and pybind docstring advertise HWC, 1xHWC, CHW, and 1xCHW image tensors. Runtime checks accept contiguous CHW and 1xCHW, but reject ordinary contiguous HWC and 1xHWC.
The implementation squeezes rank-4 input before checking ChannelsLast; PyTorch reports that memory format as false for the resulting rank-3 tensor, so the HWC branch is effectively unreachable. Direct tests with the built binding confirmed both CHW forms pass and both HWC forms fail.
The examples in #22620 document only the verified CHW behavior. Runtime and type-stub changes are intentionally excluded from that documentation-only PR.
Proposed resolution
Choose and enforce one contract:
- Implement real HWC and 1xHWC support with shape-aware conversion, or
- Remove the unreachable HWC path and narrow the pybind docstring and
_llm_runner.pyi contract to contiguous CHW and 1xCHW.
Test plan
Add positive and negative coverage for all four layouts, both supported dtypes, invalid batch size, and noncontiguous tensors.
cc @larryliu0820 @cccclai @helunwencser @jackzhxng @digantdesai
Problem
The Python type stub and pybind docstring advertise
HWC,1xHWC,CHW, and1xCHWimage tensors. Runtime checks accept contiguousCHWand1xCHW, but reject ordinary contiguousHWCand1xHWC.The implementation squeezes rank-4 input before checking
ChannelsLast; PyTorch reports that memory format as false for the resulting rank-3 tensor, so the HWC branch is effectively unreachable. Direct tests with the built binding confirmed both CHW forms pass and both HWC forms fail.The examples in #22620 document only the verified CHW behavior. Runtime and type-stub changes are intentionally excluded from that documentation-only PR.
Proposed resolution
Choose and enforce one contract:
_llm_runner.pyicontract to contiguous CHW and 1xCHW.Test plan
Add positive and negative coverage for all four layouts, both supported dtypes, invalid batch size, and noncontiguous tensors.
cc @larryliu0820 @cccclai @helunwencser @jackzhxng @digantdesai