Feat/where

docs/SUMMARY.md

@@ -89,6 +89,7 @@
     * [tensor.clip](framework/operators/tensor/tensor.clip.md)
     * [tensor.identity](framework/operators/tensor/tensor.identity.md)
     * [tensor.and](framework/operators/tensor/tensor.and.md)
+    * [tensor.where](framework/operators/tensor/tensor.where.md)
   * [Neural Network](framework/operators/neural-network/README.md)
     * [nn.relu](framework/operators/neural-network/nn.relu.md)
     * [nn.leaky\_relu](framework/operators/neural-network/nn.leaky\_relu.md)

docs/framework/compatibility.md

@@ -66,5 +66,6 @@ You can see below the list of current supported ONNX Operators:
 |                [Xor](operators/tensor/tensor.xor.md)                | :white\_check\_mark: |
 |                 [Or](operators/tensor/tensor.or.md)                 | :white\_check\_mark: |
 |             [Gemm](operators/neural-network/nn.gemm.md)             | :white\_check\_mark: |
+|             [Where](operators/tensor/tensor.where.md)               | :white\_check\_mark: |
 
-Current Operators support: **60/156 (38%)**
+Current Operators support: **61/156 (39%)**

docs/framework/operators/tensor/README.md

@@ -85,6 +85,7 @@ use orion::operators::tensor::TensorTrait;
 | [`tensor.clip`](tensor.clip.md) | Clip operator limits the given input within an interval. |
 | [`tensor.and`](tensor.and.md) | Computes the logical AND of two tensors element-wise.  |
 | [`tensor.identity`](tensor.identity.md) | Return a Tensor with the same shape and contents as input. |
+| [`tensor.where`](tensor.where.md) | Return elements chosen from x or y depending on condition. |
 
 ## Arithmetic Operations
 

docs/framework/operators/tensor/tensor.where.md

@@ -0,0 +1,48 @@
+#tensor.where
+
+```rust
+    fn where(self: @Tensor<T>, x: @Tensor<T>, y: @Tensor<T>) -> Tensor<T>;
+```
+
+Computes a new tensor by selecting values from tensor x (resp. y) at
+indices where the condition is 1 (resp. 0).
+
+## Args
+
+* `self`(`@Tensor<T>`) - The condition tensor
+* `x`(`@Tensor<T>`) - The first input tensor
+* `y`(`@Tensor<T>`) - The second input tensor
+
+## Panics
+
+* Panics if the shapes are not equal or broadcastable
+
+## Returns
+
+Return a new `Tensor<T>` of the same shape as the input with elements 
+chosen from x or y depending on the condition.
+
+## Example
+
+```rust
+use array::{ArrayTrait, SpanTrait};
+
+use orion::operators::tensor::{TensorTrait, Tensor, U32Tensor};
+
+fn where_example() -> Tensor<u32> {
+    let tensor_cond = TensorTrait::<u32>::new(
+        shape: array![2, 2].span(), data: array![0, 1, 0, 1].span(),
+    );
+
+    let tensor_x = TensorTrait::<u32>::new(
+        shape: array![2, 2].span(), data: array![2, 4, 6, 8].span(),
+    );
+
+    let tensor_y = TensorTrait::<u32>::new(
+        shape: array![2, 2].span(), data: array![1, 3, 5, 9].span(),
+    );
+
+    return tensor_cond.where(@tensor_1, @tensor_2);
+}
+>>> [1,4,5,8]
+```

nodegen/node/where.py

@@ -0,0 +1,207 @@
+import numpy as np
+from nodegen.node import RunAll
+from ..helpers import make_node, make_test, to_fp, Tensor, Dtype, FixedImpl
+
+
+class Where(RunAll):
+
+    @staticmethod
+    def where_u32():
+        def default():
+            cond = np.random.choice([1, 0], (3, 3, 3)).astype(np.uint32)
+            x = np.random.randint(0, 6, (3, 3, 3)).astype(np.uint32)
+            y = np.random.randint(0, 6, (3, 3, 3)).astype(np.uint32)
+
+            z = np.where(cond, x, y).astype(x.dtype)
+
+            cond = Tensor(Dtype.U32, cond.shape, cond.flatten())
+            x = Tensor(Dtype.U32, x.shape, x.flatten())
+            y = Tensor(Dtype.U32, y.shape, y.flatten())
+            z = Tensor(Dtype.U32, z.shape, z.flatten())
+
+            name = "where_u32"
+            make_node([cond, x, y], [z], name)
+            make_test([cond, x, y], z, "input_0.where(@input_1,@input_2)", name)
+
+        def broadcast():
+            cond = np.random.choice([1, 0], (1, 1)).astype(np.uint32)
+            x = np.random.randint(0, 6, (2, 2)).astype(np.uint32)
+            y = np.random.randint(0, 6, (1, 2)).astype(np.uint32)
+
+            z = np.where(cond, x, y).astype(x.dtype)
+
+            cond = Tensor(Dtype.U32, cond.shape, cond.flatten())
+            x = Tensor(Dtype.U32, x.shape, x.flatten())
+            y = Tensor(Dtype.U32, y.shape, y.flatten())
+            z = Tensor(Dtype.U32, z.shape, z.flatten())
+
+            name = "where_u32_broadcast"
+            make_node([cond, x, y], [z], name)
+            make_test([cond, x, y], z, "input_0.where(@input_1,@input_2)", name)
+
+        default()
+        broadcast()
+
+    @staticmethod
+    def where_i32():
+        def default():
+            cond = np.random.choice([1, 0], (3, 3, 3)).astype(np.int32)
+            x = np.random.randint(0, 6, (3, 3, 3)).astype(np.int32)
+            y = np.random.randint(0, 6, (3, 3, 3)).astype(np.int32)
+
+            z = np.where(cond, x, y).astype(x.dtype)
+
+            cond = Tensor(Dtype.I32, cond.shape, cond.flatten())
+            x = Tensor(Dtype.I32, x.shape, x.flatten())
+            y = Tensor(Dtype.I32, y.shape, y.flatten())
+            z = Tensor(Dtype.I32, z.shape, z.flatten())
+
+            name = "where_i32"
+            make_node([cond, x, y], [z], name)
+            make_test([cond, x, y], z, "input_0.where(@input_1,@input_2)", name)
+
+        def broadcast():
+            cond = np.random.choice([1, 0], (1, 1)).astype(np.int32)
+            x = np.random.randint(0, 6, (2, 2)).astype(np.int32)
+            y = np.random.randint(0, 6, (1, 2)).astype(np.int32)
+
+            z = np.where(cond, x, y).astype(x.dtype)
+
+            cond = Tensor(Dtype.I32, cond.shape, cond.flatten())
+            x = Tensor(Dtype.I32, x.shape, x.flatten())
+            y = Tensor(Dtype.I32, y.shape, y.flatten())
+            z = Tensor(Dtype.I32, z.shape, z.flatten())
+
+            name = "where_i32_broadcast"
+            make_node([cond, x, y], [z], name)
+            make_test([cond, x, y], z, "input_0.where(@input_1,@input_2)", name)
+
+        default()
+        broadcast()
+
+    @staticmethod
+    def where_i8():
+        def default():
+            cond = np.random.choice([1, 0], (3, 3, 3)).astype(np.int8)
+            x = np.random.randint(0, 6, (3, 3, 3)).astype(np.int8)
+            y = np.random.randint(0, 6, (3, 3, 3)).astype(np.int8)
+
+            z = np.where(cond, x, y).astype(x.dtype)
+
+            cond = Tensor(Dtype.I8, cond.shape, cond.flatten())
+            x = Tensor(Dtype.I8, x.shape, x.flatten())
+            y = Tensor(Dtype.I8, y.shape, y.flatten())
+            z = Tensor(Dtype.I8, z.shape, z.flatten())
+
+            name = "where_i8"
+            make_node([cond, x, y], [z], name)
+            make_test([cond, x, y], z, "input_0.where(@input_1,@input_2)", name)
+
+        def broadcast():
+            cond = np.random.choice([1, 0], (1, 1)).astype(np.int8)
+            x = np.random.randint(0, 6, (2, 2)).astype(np.int8)
+            y = np.random.randint(0, 6, (1, 2)).astype(np.int8)
+
+            z = np.where(cond, x, y).astype(x.dtype)
+
+            cond = Tensor(Dtype.I8, cond.shape, cond.flatten())
+            x = Tensor(Dtype.I8, x.shape, x.flatten())
+            y = Tensor(Dtype.I8, y.shape, y.flatten())
+            z = Tensor(Dtype.I8, z.shape, z.flatten())
+
+            name = "where_i8_broadcast"
+            make_node([cond, x, y], [z], name)
+            make_test([cond, x, y], z, "input_0.where(@input_1,@input_2)", name)
+
+        default()
+        broadcast()
+
+    @staticmethod
+    def where_fp8x23():
+        def default():
+            cond = np.random.choice([1, 0], (3, 3, 3)).astype(np.float64)
+            x = np.random.randint(0, 6, (3, 3, 3)).astype(np.float64)
+            y = np.random.randint(0, 6, (3, 3, 3)).astype(np.float64)
+
+            z = np.where(cond, x, y).astype(x.dtype)
+
+            cond = Tensor(Dtype.FP8x23, cond.shape, to_fp(
+                cond.flatten(), FixedImpl.FP8x23))
+            x = Tensor(Dtype.FP8x23, x.shape, to_fp(
+                x.flatten(), FixedImpl.FP8x23))
+            y = Tensor(Dtype.FP8x23, y.shape,  to_fp(
+                y.flatten(), FixedImpl.FP8x23))
+            z = Tensor(Dtype.FP8x23, z.shape,  to_fp(
+                z.flatten(), FixedImpl.FP8x23))
+
+            name = "where_fp8x23"
+            make_node([cond, x, y], [z], name)
+            make_test([cond, x, y], z, "input_0.where(@input_1,@input_2)", name)
+
+        def broadcast():
+            cond = np.random.choice([1, 0], (1, 1)).astype(np.float64)
+            x = np.random.randint(0, 6, (2, 2)).astype(np.float64)
+            y = np.random.randint(0, 6, (1, 2)).astype(np.float64)
+
+            z = np.where(cond, x, y).astype(x.dtype)
+
+            cond = Tensor(Dtype.FP8x23, cond.shape, to_fp(
+                cond.flatten(), FixedImpl.FP8x23))
+            x = Tensor(Dtype.FP8x23, x.shape, to_fp(
+                x.flatten(), FixedImpl.FP8x23))
+            y = Tensor(Dtype.FP8x23, y.shape,  to_fp(
+                y.flatten(), FixedImpl.FP8x23))
+            z = Tensor(Dtype.FP8x23, z.shape,  to_fp(
+                z.flatten(), FixedImpl.FP8x23))
+
+            name = "where_fp8x23_broadcast"
+            make_node([cond, x, y], [z], name)
+            make_test([cond, x, y], z, "input_0.where(@input_1,@input_2)", name)
+
+        default()
+        broadcast()
+
+    @staticmethod
+    def where_fp16x16():
+        def default():
+            cond = np.random.choice([1, 0], (3, 3, 3)).astype(np.float64)
+            x = np.random.randint(0, 6, (3, 3, 3)).astype(np.float64)
+            y = np.random.randint(0, 6, (3, 3, 3)).astype(np.float64)
+
+            z = np.where(cond, x, y).astype(x.dtype)
+
+            cond = Tensor(Dtype.FP16x16, cond.shape, to_fp(
+                cond.flatten(), FixedImpl.FP16x16))
+            x = Tensor(Dtype.FP16x16, x.shape, to_fp(
+                x.flatten(), FixedImpl.FP16x16))
+            y = Tensor(Dtype.FP16x16, y.shape,  to_fp(
+                y.flatten(), FixedImpl.FP16x16))
+            z = Tensor(Dtype.FP16x16, z.shape,  to_fp(
+                z.flatten(), FixedImpl.FP16x16))
+
+            name = "where_fp16x16"
+            make_node([cond, x, y], [z], name)
+            make_test([cond, x, y], z, "input_0.where(@input_1,@input_2)", name)
+
+        def broadcast():
+            cond = np.random.choice([1, 0], (1, 1)).astype(np.float64)
+            x = np.random.randint(0, 6, (2, 2)).astype(np.float64)
+            y = np.random.randint(0, 6, (1, 2)).astype(np.float64)
+
+            z = np.where(cond, x, y).astype(x.dtype)
+
+            cond = Tensor(Dtype.FP16x16, cond.shape, to_fp(
+                cond.flatten(), FixedImpl.FP16x16))
+            x = Tensor(Dtype.FP16x16, x.shape, to_fp(
+                x.flatten(), FixedImpl.FP16x16))
+            y = Tensor(Dtype.FP16x16, y.shape,  to_fp(
+                y.flatten(), FixedImpl.FP16x16))
+            z = Tensor(Dtype.FP16x16, z.shape,  to_fp(
+                z.flatten(), FixedImpl.FP16x16))
+
+            name = "where_fp16x16_broadcast"
+            make_node([cond, x, y], [z], name)
+            make_test([cond, x, y], z, "input_0.where(@input_1,@input_2)", name)
+
+        default()
+        broadcast()