Skip to content

Commit caf0d8d

Browse files
Make the output of ismember have the same shape as its first input (#657)
Also add documentation for the `ismember` function.
1 parent 35adfc1 commit caf0d8d

File tree

1 file changed

+22
-0
lines changed

1 file changed

+22
-0
lines changed

src/arraymancer/tensor/algorithms.nim

+22
Original file line numberDiff line numberDiff line change
@@ -255,6 +255,28 @@ proc contains*[T](t: Tensor[T], item: T): bool {.inline.}=
255255
return find(t, item) >= 0
256256

257257
proc ismember*[T](t1, t2: Tensor[T]): Tensor[bool] {.noinit.} =
258+
## Element-wise check whether elements of a Tensor are contained in another Tensor
259+
##
260+
## Inputs:
261+
## - t1: Tensor whose elements will be looked for, one by one, in `t2`
262+
## - t2: Tensor in which elements of `t1` will be looked for
263+
##
264+
## Result:
265+
## - A boolean tensor of the same shape as `t1`.
266+
## Each element indicates if `t2` contains the `t1` element that
267+
## is found in that particular position.
268+
##
269+
## Example:
270+
## ```nim
271+
## let t1 = arange(6).reshape(2, 3)
272+
## let t2 = [-3, 0, 2, 5].toTensor
273+
##
274+
## echo t1.ismember(t2)
275+
## # Tensor[system.bool] of shape "[2, 3]" on backend "Cpu"
276+
## # |true false true|
277+
## # |false false true|
278+
## ```
258279
result = newTensor[bool](t1.len)
259280
for n, it in t1.enumerate():
260281
result[n] = it in t2
282+
result = result.reshape(t1.shape)

0 commit comments

Comments
 (0)