torch.zero_ in compile

[arabel1a]

Hi! I have encountered quite a weird behaviour of the torch_npu. Consider this simple function:

def test(outputs):
    outputs.zero_()

torch._dynamo.reset()
fn = torch.compile(test)

If I simply call e.g.

fn(torch.tensor([1., 1.]).npu())
fn(torch.tensor([1., ]).npu())

Everything will work as expected. If I instead run

fn(torch.tensor([1., ]).npu())
fn(torch.tensor([1., 1.]).npu())

It will choke on the second call with this error message:

File /usr/local/python3.11.14/lib/python3.11/site-packages/torch/_inductor/utils.py:2962, in align_inputs_from_check_idxs.<locals>.run(new_inputs)
   2958 def run(new_inputs: list[InputType]) -> Any:
   2959     old_tensors, new_tensors = copy_misaligned_inputs(
   2960         new_inputs, inputs_to_check, mutated_input_idxs
   2961     )
-> 2962     out = model(new_inputs)
   2964     # If a mutated tensor was cloned to be aligned, we need to reflect back the mutation to the
   2965     # original tensor.
   2966     if len(old_tensors):

File /tmp/torchinductor_root/cb/ccbzs4bo35bo3kkkygyqfvmlknqifvwiip4zqpxlmat7ijmzfls3.py:103, in Runner.call(self, args)
    101 args.clear()
    102 s54 = arg0_1
--> 103 buf0 = empty_strided((s54, ), (1, ), device='npu', dtype=torch.float32)
    104 # Topologically Sorted Source Nodes: [zero_], Original ATen: [aten.zero]
    105 stream0 = get_raw_stream(0)

BackendCompilerFailed: backend='inductor' raised:
RuntimeError: build/CMakeFiles/torch_npu.dir/compiler_depend.ts:3950: SymIntArrayRef expected to contain only concrete integers

This error is not related simpy to broadcasting, because

# this works fine
fn(torch.tensor([1., 1.]).npu())
fn(torch.tensor([1., ]).npu())

# this fails with the same error
fn(torch.tensor([1., 1., 1., 1.]).npu())
fn(torch.tensor([1., 1. ]).npu())

# this also fails
fn(torch.tensor([1., 1. ]).npu())
fn(torch.tensor([1., 1., 1., 1.]).npu())

Using copy_() instead of zero_() does not cause this error, so it is function-specific. Interestingly, it tries to allocate a new tensor (why should zero_() do that??), and according to the error message, it is very likely that the same SymIntArrayRef error will happen with any stuff that allocates new tensors with symbolic shapes.

Environment

torch 2.9.0+cpu
torch_npu 2.9.0
CANN 8.5
Ascend 910B4