Repository navigation
local_slice_read_of_write fails on slice reads of a multi-axis advanced write #2496
Description
The local_slice_read_of_write rewrite fails with BUG IN FGRAPH.REPLACE when a graph slices a tensor written with advanced indexing on two or more axes. The rewrite turns each read slice into its own arange, so the two aranges index pointwise and read the diagonal instead of the block. The type check rejects the lower-rank replacement, so the output stays correct while every compile logs the traceback.
import numpy as np
import pytensor
import pytensor.tensor as pt
x = pt.matrix("x", shape=(4, 4))
v = pt.vector("v", shape=(2,))
written = pt.set_subtensor(x[[2, 3], [2, 3]], v)
read = written[:2, :2]
f = pytensor.function([x, v], read) # logs "Rewrite failure due to: local_slice_read_of_write"
x_value = np.arange(16.0).reshape(4, 4)
np.testing.assert_allclose(f(x_value, np.ones(2)), x_value[:2, :2]) # passeslog
<<!! BUG IN FGRAPH.REPLACE OR A LISTENER !!>> <class 'TypeError'> Cannot convert Type Vector(float64, shape=(2,)) (of Variable AdvancedSubtensor{idx_list=(0, 1)}.0) into Type Matrix(float64, shape=(2, 2)). You can try to manually convert AdvancedSubtensor{idx_list=(0, 1)}.0 into a Matrix(float64, shape=(2, 2)). local_slice_read_of_write
ERROR:pytensor.graph.rewriting.basic:Rewrite failure due to: local_slice_read_of_write
ERROR:pytensor.graph.rewriting.basic:node: Subtensor{:stop, :stop}(AdvancedSetSubtensor.0, 2, 2)
Activity