diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py index e9a985e5b728c7..1d0252018b6de7 100644 --- a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py +++ b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py @@ -676,6 +676,21 @@ def BINARY_OP(self, instr: Instruction): def BINARY_SUBSCR(self, instr: Instruction): key = self.stack.pop() container = self.stack.pop() + self.binary_subscr_operation(key, container, instr.opname) + + @call_break_graph_decorator(push_n=1) + def BINARY_SLICE(self, instr: Instruction): + end = self.stack.pop() + start = self.stack.pop() + container = self.stack.pop() + key = SliceVariable( + slice(start, end), + graph=self._graph, + tracker=DummyTracker([start, end]), + ) + self.binary_subscr_operation(key, container, instr.opname) + + def binary_subscr_operation(self, key, container, opname): assert isinstance(key, VariableBase) # TODO(xiongkun): getitem / getattr support key and attr as variable. if isinstance(key, TensorVariable) and isinstance( @@ -690,7 +705,7 @@ def BINARY_SUBSCR(self, instr: Instruction): if isinstance(key, TensorVariable): raise BreakGraphError( - f"Key is a TensorVariable in BINARY_SUBSCR, {container}[{key}]" + f"Key is a TensorVariable in {opname}, {container}[{key}]" ) result = BuiltinVariable( @@ -864,11 +879,28 @@ def STORE_SUBSCR(self, instr: Instruction): key = self.stack.pop() container = self.stack.pop() value = self.stack.pop() + self.store_subscr_operation(key, container, value, instr.opname) + + @call_break_graph_decorator(push_n=0) + def STORE_SLICE(self, instr: Instruction): + end = self.stack.pop() + start = self.stack.pop() + container = self.stack.pop() + value = self.stack.pop() + + key = SliceVariable( + slice(start, end), + graph=self._graph, + tracker=DummyTracker([start, end]), + ) + self.store_subscr_operation(key, container, value, instr.opname) + + def store_subscr_operation(self, key, container, value, opname): assert isinstance(key, VariableBase) self._graph.add_global_guarded_variable(key) if isinstance(key, TensorVariable): raise BreakGraphError( - f"Key is a TensorVariable in STORE_SUBSCR, {container}[{key}] = {value}" + f"Key is a TensorVariable in {opname}, {container}[{key}] = {value}" ) # TODO(xiongkun): support tensor[tensor] = tensor, dy2static is not the same with dygraph. container[key.get_py_value()] = value diff --git a/test/sot/skip_files_py312 b/test/sot/skip_files_py312 index 815f3a9e68b498..59cd1a37055f4e 100644 --- a/test/sot/skip_files_py312 +++ b/test/sot/skip_files_py312 @@ -2,7 +2,6 @@ ./test_11_jumps.py ./test_12_for_loop.py ./test_14_operators.py -./test_15_slice.py ./test_21_global.py ./test_analysis_inputs.py ./test_break_graph.py