Optimize slices of arange

Description

In #1429 we use the following expression to lower some xtensor indexing operations to tensor operations:

adv_idx_equivalent = arange(x_shape[i])[to_basic_idx(idx)]

This creates an unnecessary intermediate tensor and indexing operation. We should add rewrites to:

1. Convert arange(...)[slice(start, stop, step)] directly to an equivalent arange(...)
2. Convert arange(...)[scalar_index] to just the scalar value

In both cases, we need to take care of potentially negative indices (or None in slices). All variables in arange and the indexing may also be symbolic. However, we don't want to risk creating crazy graphs like #112

These optimizations would benefit both tensor and xtensor code, as the tensor rewrites would automatically apply after xtensor lowering.

Files to change

• pytensor/tensor/rewriting/subtensor_lift.py: Add new optimization patterns

Implementation details

Two new pattern rewrites should be added:

1. A rewrite that converts Subtensor(ARange, slice(...)) to a direct ARange with adjusted parameters
2. A rewrite that converts Subtensor(ARange, scalar) to a constant or scalar value

Current behavior

If we profile or print the current graph for a function like:

import pytensor
import pytensor.tensor as pt

shape_val = pt.lscalar("shape_val")
intermediate = pt.arange(shape_val)[:1]
fn = pytensor.function([shape_val], intermediate)
fn.dprint()
# Subtensor{:stop} [id A] 1
#  ├─ ARange{dtype='int64'} [id B] 0
#  │  ├─ 0 [id C]
#  │  ├─ shape_val [id D]
#  │  └─ 1 [id E]
#  └─ 1 [id F]

We'll see that the compiled graph includes:

1. An ARange operation to create a sequence from 0 to shape_val
2. A Subtensor operation to apply the slice that keeps only one value (in this case)
3. This creates an unnecessary intermediate tensor

Start in pytensor/tensor/rewriting/subtensor_lift.py and inspect the existing optimization patterns for Subtensor and ARange. Use the provided arange(shape_val)[:1] example to inspect the graph before and after rewriting. Done means slice and scalar indexing have direct ARange or scalar results while handling negative indices, None, and symbolic values without creating overly complex graphs.