Stop subtensor merge bonanza

[ricardoV94]

Summary

Closes #112
Closes #1283

local_subtensor_merge no longer expands into combinatorial switch/min/max trees on Scan outputs with symbolic shapes, and scan_save_mem still trims buffers all the way down to taps + 1 whenever a chain of constant-bound subtensors ends in a constant scalar — including when intermediate slices sit between the strip and the terminal index.

What changed

1. Gate local_subtensor_merge + add a shape-free fast path. Only call the symbolic merge_two_slices when slice bounds (and shapes, for slice+scalar) are constant. A new _safe_slice_slice_merge handles the common slice-on-slice cases without consulting shape: forward × forward (sign-aware bound combination, e.g. x[1:-1][1:-1] → x[2:-2]), x[a:b][::-1], x[::-1][a:b], x[::-1][a:b:-1], and x[a:b:-1][::-1] (last restricted to non-negative bounds). Aggressive scalar-into-slice merging moves to local_subtensor_merge_unsafe (shape_unsafe).
2. Replace while_scan_merge_subtensor_last_element with scan_merge_subtensor_chain — walks arbitrary chains of constant-bound Subtensor clients (steps ∈ {None, 1, -1}) on a Scan output, folds them into a direct raw_out[final_idx] with an Assert(n_steps >= k) for safety. Concrete simulation for constant n_steps; affine bookkeeping over raw_length for symbolic.
3. Add scan_reduce_nsteps — when every client of a Scan output is a constant scalar index, reduce n_steps to k + 1 - init_l and convert clients to negative form so save_mem can finish the job.
4. Refactor scan_save_mem — drops get_canonical_form_slice; reads buffer requirements straight off the (now-folded) negative indices. Caps prealloc extra_size at n_steps to avoid uninitialized slots.
5. Fix infer_shape and local_subtensor_shape_constant for 0-d outputs that arise after save_mem + sit_sot_to_untraced collapse.

Invariant honored

result[...][-1] (static [-1] through any chain of constant-bound, ±1-step slices) always reduces to a unit buffer, with both constant and symbolic n_steps.

Verified

Pattern	Static n=10	Symbolic n
result[-1]	0-d scan output ✓	buffer ≤ 2 + Assert(n≥1) ✓
result[-3:][-1]	0-d ✓	buffer ≤ 2 + Assert(n≥3) ✓
result[3:][-1]	0-d ✓	buffer ≤ 2 + Assert(n≥4) ✓
result[::-1][0]	0-d ✓	buffer ≤ 2 ✓
result[5::-1][-1]	0-d ✓	buffer ≤ 2 ✓
result[:-1][::-1][-1]	0-d ✓	reduced ✓

tests/scan/test_rewriting.py: 59 passed / 2 skipped. tests/tensor/rewriting/test_subtensor.py::TestLocalSubtensorMerge: 13 passed / 2 skipped.

Benchmarks

Compile-time and post-rewrite node counts at the pre-gate baseline (d35fb51b1) vs this branch, mode excluding("fusion") so the rewriter output is visible. Composite fusion otherwise lumps the symbolic switch tree into a single node, hiding the count but not the rewrite-time cost. constant_folding is not hiding anything — the switches reference symbolic Shape_i(x) and don't fold.

Benchmark	Before	After	Speedup
x[1:-1] × 3	82 apply, 228 ms	2 apply, 27 ms	8×
x[1:-1] × 5	166 apply, 490 ms	2 apply, 45 ms	11×
x[1:-1] × 8	292 apply, 835 ms	2 apply, 74 ms	11×
grad(xs[-1], x0) symbolic-n Scan (#112's worst case)	348 apply, 1260 ms	22 apply, 114 ms	11×

Covered by tests/benchmarks/test_subtensor.py::test_local_subtensor_merge_compile_benchmark and tests/benchmarks/test_scan.py::test_scan_grad_subtensor_compile_benchmark

Limitations (deliberate)

• While-scans only fold to final_idx == -1: the only number of steps we can even be sure a while scan takes is 1 (if n_steps>0, which we assert)
• Slice steps of magnitude > 1 are not handled.
• concatenate([rev, zeros])[k] from the while-scan gradient path stays opaque to chain folding (would need Consider lifting Subtensor through Joins #919)

[ricardoV94]

Possibly fixes: #1288 need to check