Extend `local_useless_slice` to IncSubtensor and AdvancedIncSubetensor Ops

[ricardoV94]

Description

The rewrite local_useless_slice canonicalizes slicing operations so that x[0:-1:2] and x[::2] are made equivalent.

This should also apply to IncSubtensor operations with slices where x[0:-1:2].set(y) and x[::2].set(y) are made equivalent. The only difference is the branch where we find zero indices and just return x, in that case for the IncSubtensor, we should return full_like(x, y)

pytensor/pytensor/tensor/rewriting/subtensor.py

Lines 337 to 406 in b8dbd4c

	@register_infer_shape
	@register_useless
	@register_canonicalize
	@register_specialize
	@register_stabilize
	@node_rewriter([Subtensor])
	def local_useless_slice(fgraph, node):
	"""
	Remove Subtensor of the form:
	1. X[0, :] -> X[0]
	2. X[:] -> X
	Also, rewrite Subtensor of the form:
	X[0:7:1] -> X[None:None:None]
	where X is a vector of length 7
	"""
	idxs = get_idx_list(node.inputs, node.op.idx_list)
	x = node.inputs[0]
	if not idxs:
	return [node.inputs[0]]
	new_idxs = list(idxs)
	change_flag = False
	last_useful_idx = -1
	for dim, s in enumerate(new_idxs):
	if not isinstance(s, slice):
	last_useful_idx = dim
	continue
	if s == slice(None):
	continue
	start = s.start
	stop = s.stop
	step = s.step
	if (
	start is not None
	and extract_constant(start, only_process_constants=True) == 0
	):
	change_flag = True
	start = None
	if (
	stop is not None
	and x.type.shape[dim] is not None
	and extract_constant(stop, only_process_constants=True) == x.type.shape[dim]
	):
	change_flag = True
	stop = None
	if (
	step is not None
	and extract_constant(step, only_process_constants=True) == 1
	):
	change_flag = True
	step = None
	if not (start is None and stop is None and step is None):
	last_useful_idx = dim
	new_idxs[dim] = slice(start, stop, step)
	if change_flag or ((last_useful_idx + 1) < len(idxs)):
	out = x[tuple(new_idxs[: last_useful_idx + 1])]
	# Copy over previous output stacktrace
	copy_stack_trace(node.outputs, out)
	return [out]