Lower PReLU, pixel_shuffle/unshuffle, and step>1 slice to GPU-clean forms

[john-rocky]

Summary

Follow-up to #1079, same theme: several common ops decompose into primitives the GPU delegate (ML Drift / LITERT_CL) rejects, even though a GPU-clean equivalent exists. This PR fixes the three worst offenders in litert-torch's decomposition/lowering layer. Op sets read from the converted .tflite via Interpreter._get_ops_details() on main HEAD (7781284):

Op	Before (main HEAD)	After this PR	Numerics
PReLU (nn.PReLU)	GREATER + SELECT (+MUL) — both GPU-rejected	RELU/MUL/SUB	exact, max|diff| = 0.0
PixelShuffle / pixel_unshuffle	rank-6 RESHAPE (>4D cap)	rank-4 RESHAPE/TRANSPOSE	exact, max|diff| = 0.0
step>1 slice / Focus stem (x[..., ::2, ::2])	GATHER_ND — GPU-rejected †	single STRIDED_SLICE builtin	exact, max|diff| = 0.0

† JAX-version dependent — see Correction on the slice half (item 3) below.

Each of these today forces a model off the GPU or into onnx2tf / private converter patches: PReLU blocks restoration/SR nets (e.g. Real-ESRGAN), the rank-6 pixel-shuffle blocks sub-pixel upsampler heads, and GATHER_ND blocks YOLOX/YOLOv5 Focus stems and ViT patch-embed-via-slicing.

What this PR does

1. aten._prelu_kernel decomposition override (_decomp_registry.py)
The core aten decomposition is where(x > 0, x, w * x) → GREATER + SELECT. Overridden with the numerically identical (including NaN propagation) relu form relu(x) - w * relu(-x), following the existing _safe_softmax override precedent in the same file.

2. aten.pixel_shuffle / aten.pixel_unshuffle decomposition overrides (_decomp_registry.py)
The default decomposition materializes a rank-6 intermediate (reshape → 6D permute → reshape), but GPU delegates cap tensor rank at 4. The override folds batch and channel into one dimension and interleaves one spatial axis at a time, so every intermediate stays rank 4: 3 reshapes + 2 transposes, exact for any batch shape and upscale factor (also handles the 3D no-batch case).

3. aten.slice / aten.slice_copy strided lowering (_jax_lowerings/lowerings.py)
torchax implements slice with jnp basic indexing; JAX's strided path lowers to lax.gather → TFLite GATHER_ND. For fully static step > 1 slices we now emit jax.lax.slice, which the converter legalizes to a single STRIDED_SLICE. Step-1 and dynamic slices keep the torchax lowering unchanged, so generative-path slicing (masks, KV caches, dynamic dims) is unaffected.

Verification

• Op-level numerics (convert → ai_edge_litert Interpreter → compare vs eager): all new cases exact (max|diff| = 0.0), max tensor rank ≤ 4, no GATHER_ND/GREATER/SELECT. Covers: pixel_shuffle r∈{1,2,3} incl. 3D input, pixel_unshuffle r∈{2,3}, PReLU per-channel + scalar alpha, strided slice with negative/None/sys.maxsize bounds, and step-1 regression cases.
• YOLOX Focus stem (cat([x[...,::2,::2], x[...,1::2,::2], x[...,::2,1::2], x[...,1::2,1::2]], 1)): now CONCATENATION + STRIDED_SLICE only (was GATHER_ND).
• Generative regression: export_hf end-to-end on a small Llama; float CPU parity vs HF reference corr = 1.000000, top-1 100%, max|diff| = 7.5e-08. KV-cache tests green.
• New cases added to test_core_aten_ops.py (pixel_shuffle/unshuffle, prelu, strided slice/slice_copy).

On-device measurement (Pixel 8a, LITERT_CL / ML Drift)

I converted the Real-ESRGAN general-x4v3 super-resolution model twice — once on pristine
main HEAD (b66af07), once with this PR — and loaded each on a Pixel 8a through the LiteRT
CompiledModel GPU path. It is SRVGGNetCompact with the official weights: 33 × nn.PReLU and
nn.PixelShuffle(4), stock modules, no model-side rewrite. The two .tflite files are
identical except for the litert-torch version.

litert-torch	GPU delegate (LITERT_CL)	CompiledModel(Accelerator.GPU)	latency (GPU+CPU)
main HEAD (before)	Replacing 134 out of 141 node(s) … 3 partitions	Create fails	69.2 ms
this PR (after)	Replacing 176 out of 176 node(s) … 1 partition	GPU ready	35.7 ms

Before, the delegate rejects the pixel-shuffle tail outright:

E tflite : Following operations are not supported by GPU delegate:
E tflite : RESHAPE: Tensor "…PixelShuffle_upsampler;3" has bad input dims size: 6.
E tflite : RESHAPE: Tensor "…PixelShuffle_upsampler;4" has bad input dims size: 6.
E tflite : TRANSPOSE: Max version supported: 5. Requested version 6.
E tflite : 134 operations will run on the GPU, and the remaining 7 operations will run on the CPU.
I tflite : Replacing 134 out of 141 node(s) with delegate (LITERT_CL) node, yielding 3 partitions for subgraph 0.
I tflite : Replacing 7 out of 8 node(s) with delegate (TfLiteXNNPackDelegate) node, yielding 3 partitions for subgraph 0.

With the GPU accelerator alone and no CPU fallback, CompiledModel::Create returns an error and
the model never loads. With CPU fallback enabled it loads, but splits into 3 partitions and runs
the tail on XNNPACK. After this PR no op is rejected, the graph is one partition, and
CompiledModel(Accelerator.GPU) succeeds.

Latency is best-of-10 with the output read back inside the timed loop (enqueue + compute +
readback), CPU fallback enabled so that both sides actually run: 69.2 ms → 35.7 ms.

The result is unchanged. Each device GPU output matches its host CPU .tflite reference at
corr = 0.994894 before and 0.994892 after (fp16 GPU compute vs fp32 CPU), and the two device
outputs agree with each other at corr = 0.999998, max|diff| = 2.4e-04.

Converter-side, on the same two files: before 141 nodes, GREATER ×33 + SELECT ×33, max tensor
rank 6, 2 tensors >4D; after 176 nodes, no GPU-hostile op, max rank 4, 0 tensors >4D.
.tflite-vs-eager max|diff| is identical on both sides (4.86e-06). The 35 extra nodes are
exactly the cost of the two rewrites: +1 op per PReLU (33) and +2 for the rank-4 pixel shuffle.

This is the model in my open super-resolution sample, google-ai-edge/litert-samples#170. To reach
the GPU today, that sample's conversion script rewrites PReLU into relu(x) − a·relu(−x) and
PixelShuffle into a zero-stuffed transposed conv, by hand, model-side. With this PR the stock
graph is already GPU-clean, so neither rewrite is needed.

Correction on the slice half (item 3)

Re-measured on current main HEAD: the step>1-slice → GATHER_ND behaviour is JAX-version
dependent, which my original table (main @ 7781284) did not say.

env	x[:, :, ::2, ::2] on main b66af07	with this PR
jax 0.6.2	GATHER_ND ×2	STRIDED_SLICE ×2
jax 0.11.1	STRIDED_SLICE ×2	STRIDED_SLICE ×2 (no change)

The cause is upstream of both litert-torch and the TFLite converter — JAX's own lowering of
strided basic indexing changed:

jax.jit(lambda x: x[:, :, ::2, ::2]).lower(jnp.zeros((1, 8, 16, 16))).as_text()
# jax 0.6.2  -> stablehlo.gather (+ iota/multiply/add index math)
# jax 0.11.1 -> stablehlo.slice

requirements.txt leaves jax[cpu] unpinned, so a fresh py3.12/3.13 install lands on jax 0.11.x
and will not reproduce that row. A py3.10 install still does: 0.6.2 is the newest jax that
publishes a py3.10 wheel. The YOLOX Focus stem behaves the same way — GATHER_ND ×6 on jax 0.6.2,
STRIDED_SLICE ×6 on jax 0.11.1.

The PReLU and pixel_shuffle/unshuffle rows reproduce identically on both.

So item 3 earns its place only on the older-JAX installs; on a current JAX it is a no-op, with
identical op sets and counts before and after. I'm happy to split it into its own PR, or drop it,
if you'd rather keep this one to the two decomposition overrides.

Notes / follow-up

The native PRELU and DEPTH_TO_SPACE/SPACE_TO_DEPTH TFLite builtins would be even better targets, but emitting them needs converter-side pattern or composite support (no stablehlo→TFL legalization currently produces them) — out of litert-torch's decomposition layer, so this PR takes the ops to GPU-clean ≤4D elementwise/reshape forms instead, and STRIDED_SLICE (where a direct legalization does exist) to the native builtin.

Together with #1079 (attention view chains within rank 4), this clears the decomposition-side GPU blockers for the common vision-model patterns.

[outtanames]

Nice work! Just a handful of nits but overall g2g.

[john-rocky]

Thanks for the review! All four comments are addressed in 54c5137 — details in the inline replies. Summary: the pixel_shuffle JAX lowering was indeed dead and is deleted; the slice staticness check now treats export_utils.IR_DYNAMIC as dynamic (and requires the whole shape to be static); and the tests now cover 3D (no-batch) pixel_shuffle/unshuffle, sys.maxsize slice bounds, plus a regression test asserting no GATHER_ND/GREATER/SELECT-legalizing ops and max rank ≤ 4.

[outtanames]

@john-rocky mind taking care of the conflicts on _decomp_registry.py

[john-rocky]

Done — rebased onto main. The _decomp_registry.py conflict was both branches appending new decomps (main's 1D-conv promotion and this PR's PReLU / pixel_shuffle), so both are kept. The test_gpu_clean_lowering regression tests pass on the rebased branch.

[john-rocky]

Rebased on main and added a from-scratch on-device A/B (Pixel 8a, LITERT_CL / ML Drift GPU) to the description.

I converted the Real-ESRGAN general-x4v3 image model twice — stock nn.PReLU and nn.PixelShuffle, no model-side rewrite — on pristine main HEAD (b66af07) vs this branch, so the only difference is the litert-torch version, then loaded each via CompiledModel(Accelerator.GPU):

• before (main): Replacing 134 out of 141 node(s) … 3 partitions, with RESHAPE … has bad input dims size: 6 on the pixel-shuffle tail → CompiledModel::Create fails on the GPU-only path. With CPU fallback allowed it loads as 3 partitions with a CPU tail, 69.2 ms.
• after (this PR): Replacing 176 out of 176 node(s) … 1 partition → GPU ready, 35.7 ms, and the device output still matches the CPU reference (corr 0.9949, unchanged).

One correction in the same edit. Re-measured on current main, the third item (step>1 slice → GATHER_ND) is JAX-version dependent: jax 0.6.2 lowers strided basic indexing to stablehlo.gather, jax 0.11.1 to stablehlo.slice, so a fresh py3.12 install no longer reproduces that row. There is a two-line repro in the description. I'm happy to split that item into its own PR or drop it, if you'd rather keep this one to the two decomposition overrides.