[REQ] Add a wp.map() host-side API to apply a function to each element of an array #694
Closed
Description
Description
Applies a function to elements of a Warp array.
def map(
func: Callable | wp.Function,
*inputs: wp.array | Any,
out: wp.array | tuple[wp.array] | None = None,
return_kernel: bool = False,
block_dim=256,
device: Devicelike = None,
) -> wp.array | tuple[wp.array] | wp.Kernel:a = wp.array([1, 2, 3], dtype=wp.float32)
b = wp.array([4, 5, 6], dtype=wp.float32)
c = wp.array([7, 8, 9], dtype=wp.float32)
result = wp.map(lambda x, y, z: x + 2.0 * y - z, a, b, c)
print(result.numpy()) # Output: [2. 4. 6.]xs = wp.array([-1.0, 0.0, 1.0], dtype=wp.float32)
wp.map(wp.clamp, xs, -0.5, 0.5, out=xs)
print(xs.numpy()) # Output: [-0.5 0. 0.5]Additionally, built-in operators, like +, -, *, /, etc. will be allowed to work directly from the Python context as vectorized operations on Warp arrays, similar to how users are used to tensor operations in PyTorch or JAX:
a = wp.array([1, 2, 3], dtype=wp.float32)
b = wp.array([4, 5, 6], dtype=wp.float32)
print((a + b).numpy()) # Output: [5, 7, 9]