Tuple-unpack assignment to wp.ref parameters fails

[shi-eric]

Description

Tuple-unpack assignment to wp.ref[T] parameters does not behave like sequential assignment.

wp.ref[T] assignment is documented as mutating the caller's storage, and simple assignment is covered by test_ref_simple_assignment_mutates. Tuple-unpack assignment should preserve the same behavior for each ref target.

Reproducer

import warp as wp


@wp.func
def set_pair(x: wp.ref[wp.float32], y: wp.ref[wp.float32], a: wp.float32, b: wp.float32):
    x, y = a, b


@wp.kernel(enable_backward=False)
def kernel(out: wp.array(dtype=wp.float32)):
    x = wp.float32(0.0)
    y = wp.float32(0.0)

    set_pair(x, y, wp.float32(3.0), wp.float32(4.0))

    out[0] = x
    out[1] = y


out = wp.zeros(2, dtype=wp.float32, device="cpu")
wp.launch(kernel, dim=1, outputs=[out], device="cpu")
print(out.numpy())

Expected behavior

The tuple assignment should be equivalent to:

x = a
y = b

and out should contain:

[3.0, 4.0]

Actual behavior

Codegen rejects the assignment:

WarpCodegenTypeError: Error, assigning to existing symbol x (wp.ref[float32]) with different type (<class 'warp._src.types.float32'>)

Verification

Reproduced on main at 57f7a778c8cb08c73b3e6ed820759f9ee26ad10f with a CPU launch.

Related

Follow-up from #1277.

[c0d1f1ed]

Fix merged in 4fbd135.

[shi-eric]

This is addressed by 4fbd135.

The change makes tuple-unpack assignment to wp.ref[T] parameters store
through each reference instead of rebinding the parameter. The reported
x, y = a, b case now mutates caller storage as documented. Reference-valued
RHS expressions are snapshotted before stores, preserving simultaneous
assignment for swaps and mixed ref/local targets.

Public API behavior:

• In Warp kernel scope, existing wp.ref[T] parameters on @wp.func now
  support tuple-unpack assignment. No host-side Python API was added.

Test coverage:

• New tests in warp/tests/test_ref.py cover the reported assignment,
  reference swaps, mixed ref/local targets, and preservation of gradients
  for ordinary array-element tuple unpack.