From b2620e6279a16cf514c4bbe67e5e2f34f1d71291 Mon Sep 17 00:00:00 2001 From: Mamy Ratsimbazafy Date: Mon, 20 Apr 2020 01:25:32 +0200 Subject: [PATCH] Implement Numpy fancy indexing (#434) * index_select should use SomeInteger not SOmeNumber * Overload index_select for arrays and sequences * Masked Selector overload for openarrays * Add masked overload for regular arrays and sequences * Initial support of Numpy fancy indexing: index select * Fix broadcast operators from #429 using deprecated syntax * Stash dispatcher, working with types in macros is a minefield https://github.com/nim-lang/Nim/issues/14021 * Masked indexing: closes #400, workaround https://github.com/nim-lang/Nim/issues/14021 * Test for full masked fancy indexing * Add index_fill * Tensor mutation via fancy indexing * Add tests for index mutation via fancy indexing * Fancy indexing: supports broadcasting a value to a masked assignation * Detect wrong mask or tensor axis length * masked axis assign value test * Add masked assign of broadcastable tensor * Tag for changelog [skip ci] --- changelog.md | 3 + src/arraymancer.nim | 2 +- src/nn/layers/embedding.nim | 2 +- src/nn_primitives/nnp_embedding.nim | 4 +- src/private/ast_utils.nim | 10 +- src/tensor/accessors_macros_write.nim | 8 +- src/tensor/init_cpu.nim | 2 +- src/tensor/operators_comparison.nim | 12 +- .../private/p_accessors_macros_read.nim | 135 +++++++++- .../private/p_accessors_macros_write.nim | 68 ++++- src/tensor/selectors.nim | 240 +++++++++++++++--- tests/_split_tests/tests_tensor_part01.nim | 2 + tests/manual_checks/fancy_indexing.py | 190 ++++++++++++++ tests/tensor/test_fancy_indexing.nim | 203 +++++++++++++++ tests/tensor/test_selectors.nim | 179 ++++++++++++- tests/tests_cpu.nim | 2 + 16 files changed, 993 insertions(+), 69 deletions(-) create mode 100644 tests/manual_checks/fancy_indexing.py create mode 100644 tests/tensor/test_fancy_indexing.nim diff --git a/changelog.md b/changelog.md index fbc37bf34..498b2f06c 100644 --- a/changelog.md +++ b/changelog.md @@ -1,6 +1,9 @@ Arraymancer v0.x.x ===================================================== +Changes (TODO): + - Fancy Indexing (#434) + Arraymancer v0.6.0 Jan. 09 2020 ===================================================== diff --git a/src/arraymancer.nim b/src/arraymancer.nim index 4c9603886..ed62bff2b 100644 --- a/src/arraymancer.nim +++ b/src/arraymancer.nim @@ -39,6 +39,6 @@ export tensor, einsum when not defined(no_lapack): - # THe ml module also does not export everything is LAPACK is not available + # The ml module also does not export everything is LAPACK is not available import ./linear_algebra/linear_algebra export linear_algebra diff --git a/src/nn/layers/embedding.nim b/src/nn/layers/embedding.nim index 4f51cd233..45e67bbc1 100644 --- a/src/nn/layers/embedding.nim +++ b/src/nn/layers/embedding.nim @@ -63,7 +63,7 @@ proc embedding_cache[TT, Idx]( ) -proc embedding*[TT; Idx: byte or char or SomeNumber]( +proc embedding*[TT; Idx: byte or char or SomeInteger]( input_vocab_id: Tensor[Idx], weight: Variable[TT], padding_idx: Idx = -1, diff --git a/src/nn_primitives/nnp_embedding.nim b/src/nn_primitives/nnp_embedding.nim index 5005e4caf..5a6f5631e 100644 --- a/src/nn_primitives/nnp_embedding.nim +++ b/src/nn_primitives/nnp_embedding.nim @@ -8,7 +8,7 @@ import proc flatten_idx(t: Tensor): Tensor {.inline.}= t.reshape(t.size) -proc embedding*[T; Idx: byte or char or SomeNumber]( +proc embedding*[T; Idx: byte or char or SomeInteger]( vocab_id: Tensor[Idx], weight: Tensor[T] ): Tensor[T] = @@ -54,7 +54,7 @@ proc embedding*[T; Idx: byte or char or SomeNumber]( let shape = vocab_id.shape & weight.shape[1] result = weight.index_select(0, vocab_id.flatten_idx).reshape(shape) -proc embedding_backward*[T; Idx: byte or char or SomeNumber]( +proc embedding_backward*[T; Idx: byte or char or SomeInteger]( dWeight: var Tensor[T], vocab_id: Tensor[Idx], dOutput: Tensor[T], diff --git a/src/private/ast_utils.nim b/src/private/ast_utils.nim index c5557ca2e..97548f49f 100644 --- a/src/private/ast_utils.nim +++ b/src/private/ast_utils.nim @@ -16,7 +16,6 @@ import macros - proc hasType*(x: NimNode, t: static[string]): bool {. compileTime .} = ## Compile-time type checking sameType(x, bindSym(t)) @@ -25,6 +24,15 @@ proc isInt*(x: NimNode): bool {. compileTime .} = ## Compile-time type checking hasType(x, "int") +proc isBool*(x: NimNode): bool {. compileTime .} = + ## Compile-time type checking + hasType(x, "bool") + +proc isOpenarray*(x: NimNode): bool {. compileTime .} = + ## Compile-time type checking + doAssert false, "This is broken for generics https://github.com/nim-lang/Nim/issues/14021" + hasType(x, "array") or hasType(x, "seq") or hasType(x, "openArray") + proc isAllInt*(slice_args: NimNode): bool {. compileTime .} = ## Compile-time type checking result = true diff --git a/src/tensor/accessors_macros_write.nim b/src/tensor/accessors_macros_write.nim index 6918985bf..cae31fb8a 100644 --- a/src/tensor/accessors_macros_write.nim +++ b/src/tensor/accessors_macros_write.nim @@ -52,10 +52,10 @@ macro `[]=`*[T](t: var Tensor[T], args: varargs[untyped]): untyped = slice_typed_dispatch_mut(`t`, `new_args`,`val`) -# # Linked to: https://github.com/mratsim/Arraymancer/issues/52 +# Linked to: https://github.com/mratsim/Arraymancer/issues/52 # Unfortunately enabling this breaksthe test suite # "Setting a slice from a view of the same Tensor" - +# # macro `[]`*[T](t: var AnyTensor[T], args: varargs[untyped]): untyped = # ## Slice a Tensor or a CudaTensor # ## Input: @@ -69,6 +69,6 @@ macro `[]=`*[T](t: var Tensor[T], args: varargs[untyped]): untyped = # ## For CudaTensor only, this is a no-copy operation, data is shared with the input. # ## This proc does not guarantee that a ``let`` value is immutable. # let new_args = getAST(desugar(args)) - +# # result = quote do: -# slice_typed_dispatch_var(`t`, `new_args`) \ No newline at end of file +# slice_typed_dispatch_var(`t`, `new_args`) diff --git a/src/tensor/init_cpu.nim b/src/tensor/init_cpu.nim index 49750ff70..23f6f982b 100644 --- a/src/tensor/init_cpu.nim +++ b/src/tensor/init_cpu.nim @@ -91,7 +91,7 @@ proc newTensorWith*[T](shape: MetadataArray, value: T): Tensor[T] {.noInit, noSi {.unroll: 8.} tval = value -proc toTensor*(s:openarray, dummy_bugfix: static[int] = 0 ): auto {.noSideEffect.} = +proc toTensor*(s:openarray, dummy_bugfix: static[int] = 0): auto {.noSideEffect.} = ## Convert an openarray to a Tensor ## Input: ## - An array or a seq (can be nested) diff --git a/src/tensor/operators_comparison.nim b/src/tensor/operators_comparison.nim index e55044f06..c0e42c36a 100644 --- a/src/tensor/operators_comparison.nim +++ b/src/tensor/operators_comparison.nim @@ -98,38 +98,38 @@ template gen_broadcasted_scalar_comparison(op: untyped): untyped {.dirty.} = result = map_inline(t): op(x, value) -proc `.==`*[T](t: Tensor[T], value : T): Tensor[bool] {.noInit.} = +proc `==.`*[T](t: Tensor[T], value : T): Tensor[bool] {.noInit.} = ## Tensor element-wise equality with scalar ## Returns: ## - A tensor of boolean gen_broadcasted_scalar_comparison(`==`) -proc `.!=`*[T](t: Tensor[T], value : T): Tensor[bool] {.noInit.} = +proc `!=.`*[T](t: Tensor[T], value : T): Tensor[bool] {.noInit.} = ## Tensor element-wise inequality with scalar ## Returns: ## - A tensor of boolean gen_broadcasted_scalar_comparison(`!=`) -proc `.<=`*[T](t: Tensor[T], value : T): Tensor[bool] {.noInit.} = +proc `<=.`*[T](t: Tensor[T], value : T): Tensor[bool] {.noInit.} = ## Tensor element-wise lesser or equal with scalar ## Returns: ## - A tensor of boolean gen_broadcasted_scalar_comparison(`<=`) -proc `.<`*[T](t: Tensor[T], value : T): Tensor[bool] {.noInit.} = +proc `<.`*[T](t: Tensor[T], value : T): Tensor[bool] {.noInit.} = ## Tensor element-wise lesser than a scalar ## Returns: ## - A tensor of boolean gen_broadcasted_scalar_comparison(`<`) -proc `.>=`*[T](t: Tensor[T], value : T): Tensor[bool] {.noInit.} = +proc `>=.`*[T](t: Tensor[T], value : T): Tensor[bool] {.noInit.} = ## Tensor element-wise greater or equal than a scalar ## Returns: ## - A tensor of boolean gen_broadcasted_scalar_comparison(`>=`) -proc `.>`*[T](t: Tensor[T], value : T): Tensor[bool] {.noInit.} = +proc `>.`*[T](t: Tensor[T], value : T): Tensor[bool] {.noInit.} = ## Tensor element-wise greater than a scalar ## Returns: ## - A tensor of boolean diff --git a/src/tensor/private/p_accessors_macros_read.nim b/src/tensor/private/p_accessors_macros_read.nim index 534c66a95..c069905df 100644 --- a/src/tensor/private/p_accessors_macros_read.nim +++ b/src/tensor/private/p_accessors_macros_read.nim @@ -21,6 +21,8 @@ import ../../private/ast_utils, ./p_checks, ./p_accessors, sequtils, macros +from ../init_cpu import toTensor + template slicerImpl*[T](result: AnyTensor[T]|var AnyTensor[T], slices: ArrayOfSlices): untyped = ## Slicing routine @@ -117,21 +119,148 @@ proc slicer*[T](t: Tensor[T], slices: ArrayOfSlices): Tensor[T] {.noInit,noSideE # ######################################################################### # Dispatching logic +type FancySelectorKind* = enum + FancyNone + FancyIndex + FancyMaskFull + FancyMaskAxis + # Workaround needed for https://github.com/nim-lang/Nim/issues/14021 + FancyUnknownFull + FancyUnknownAxis + +proc getFancySelector*(ast: NimNode, axis: var int, selector: var NimNode): FancySelectorKind = + ## Detect indexing in the form + ## - "tensor[_, _, [0, 1, 4], _, _] + ## - "tensor[_, _, [0, 1, 4], `...`] + ## or with the index selector being a tensor + result = FancyNone + var foundNonSpanOrEllipsis = false + var ellipsisAtStart = false + + template checkNonSpan(): untyped {.dirty.} = + doAssert not foundNonSpanOrEllipsis, + "Fancy indexing is only compatible with full spans `_` on non-indexed dimensions" & + " and/or ellipsis `...`" + + var i = 0 + while i < ast.len: + let cur = ast[i] + + # Important: sameType doesn't work for generic type like Array, Seq or Tensors ... + # https://github.com/nim-lang/Nim/issues/14021 + if cur.sameType(bindSym"SteppedSlice") or cur.isInt(): + if cur.eqIdent"Span": + discard + else: + doAssert result == FancyNone + foundNonSpanOrEllipsis = true + elif cur.sameType(bindSym"Ellipsis"): + if i == ast.len - 1: # t[t.sum(axis = 1) >. 0.5, `...`] + doAssert not ellipsisAtStart, "Cannot deduce the indexed/sliced dimensions due to ellipsis at the start and end of indexing." + ellipsisAtStart = false + elif i == 0: # t[`...`, t.sum(axis = 0) >. 0.5] + ellipsisAtStart = true + else: + # t[0 ..< 10, `...`, t.sum(axis = 0) >. 0.5] is unsupported + # so we tag as "foundNonSpanOrEllipsis" + foundNonSpanOrEllipsis = true + elif cur.kind == nnkBracket: + checkNonSpan() + axis = i + if cur[0].kind == nnkIntLit: + result = FancyIndex + selector = cur + elif cur[0].isBool(): + let full = i == 0 and ast.len == 1 + result = if full: FancyMaskFull else: FancyMaskAxis + selector = cur + else: + # byte, char, enums are all represented by integers in the VM + error "Fancy indexing is only possible with integers or booleans" + else: + checkNonSpan() + axis = i + let full = i == 0 and ast.len == 1 + result = if full: FancyUnknownFull else: FancyUnknownAxis + selector = cur + inc i + + # Handle ellipsis at the start + if result != FancyNone and ellipsisAtStart: + axis = ast.len - axis + macro slice_typed_dispatch*(t: typed, args: varargs[typed]): untyped = ## Typed macro so that isAllInt has typed context and we can dispatch. ## If args are all int, we dispatch to atIndex and return T ## Else, all ints are converted to SteppedSlices and we return a Tensor. ## Note, normal slices and `_` were already converted in the `[]` macro ## TODO in total we do 3 passes over the list of arguments :/. It is done only at compile time though + + # Point indexing + # ----------------------------------------------------------------- if isAllInt(args): - result = newCall(bindSym("atIndex"), t) + result = newCall(bindSym"atIndex", t) for slice in args: result.add(slice) - else: - result = newCall(bindSym("slicer"), t) + return + + # Fancy indexing + # ----------------------------------------------------------------- + # Cannot depend/bindSym the "selectors.nim" proc + # Due to recursive module dependencies + var selector: NimNode + var axis: int + let fancy = args.getFancySelector(axis, selector) + if fancy == FancyIndex: + return newCall( + ident"index_select", + t, newLit axis, selector + ) + if fancy == FancyMaskFull: + return newCall( + ident"masked_select", + t, selector + ) + elif fancy == FancyMaskAxis: + return newCall( + ident"masked_axis_select", + t, selector, newLit axis + ) + + # Slice indexing + # ----------------------------------------------------------------- + if fancy == FancyNone: + result = newCall(bindSym"slicer", t) for slice in args: if isInt(slice): ## Convert [10, 1..10|1] to [10..10|1, 1..10|1] result.add(infix(slice, "..", infix(slice, "|", newIntLitNode(1)))) else: result.add(slice) + return + + # Fancy bug in Nim compiler + # ----------------------------------------------------------------- + # We need to drop down to "when a is T" to infer what selector to call + # as `getType`/`getTypeInst`/`getTypeImpl`/`sameType` + # are buggy with generics + # due to https://github.com/nim-lang/Nim/issues/14021 + let lateBind_masked_select = ident"masked_select" + let lateBind_masked_axis_select = ident"masked_axis_select" + let lateBind_index_select = ident"index_select" + + result = quote do: + type FancyType = typeof(`selector`) + when FancyType is (array or seq): + type FancyTensorType = typeof(toTensor(`selector`)) + else: + type FancyTensorType = FancyType + when FancyTensorType is Tensor[bool]: + when FancySelectorKind(`fancy`) == FancyUnknownFull: + `lateBind_masked_select`(`t`, `selector`) + elif FancySelectorKind(`fancy`) == FancyUnknownAxis: + `lateBind_masked_axis_select`(`t`, `selector`, `axis`) + else: + {.error: "Unreachable".} + else: + `lateBind_index_select`(`t`, `axis`, `selector`) diff --git a/src/tensor/private/p_accessors_macros_write.nim b/src/tensor/private/p_accessors_macros_write.nim index 354a3fd0b..8a51f49a5 100644 --- a/src/tensor/private/p_accessors_macros_write.nim +++ b/src/tensor/private/p_accessors_macros_write.nim @@ -183,13 +183,46 @@ proc slicerMut*[T](t: var Tensor[T], macro slice_typed_dispatch_mut*(t: typed, args: varargs[typed], val: typed): untyped = ## Assign `val` to Tensor T at slice/position `args` + + # Point indexing + # ----------------------------------------------------------------- if isAllInt(args): - result = newCall(bindSym("atIndexMut"), t) + result = newCall(bindSym"atIndexMut", t) for slice in args: result.add(slice) result.add(val) - else: - result = newCall(bindSym("slicerMut"), t) + return + + # Fancy indexing + # ----------------------------------------------------------------- + # Cannot depend/bindSym the "selectors.nim" proc + # Due to recursive module dependencies + var selector: NimNode + var axis: int + let fancy = args.getFancySelector(axis, selector) + if fancy == FancyIndex: + return newCall( + ident"index_fill", + t, newLit axis, selector, + val + ) + if fancy == FancyMaskFull: + return newCall( + ident"masked_fill", + t, selector, + val + ) + elif fancy == FancyMaskAxis: + return newCall( + ident"masked_axis_fill", + t, selector, newLit axis, + val + ) + + # Slice indexing + # ----------------------------------------------------------------- + if fancy == FancyNone: + result = newCall(bindSym"slicerMut", t) for slice in args: if isInt(slice): ## Convert [10, 1..10|1] to [10..10|1, 1..10|1] @@ -197,10 +230,39 @@ macro slice_typed_dispatch_mut*(t: typed, args: varargs[typed], val: typed): unt else: result.add(slice) result.add(val) + return + + # Fancy bug in Nim compiler + # ----------------------------------------------------------------- + # We need to drop down to "when a is T" to infer what selector to call + # as `getType`/`getTypeInst`/`getTypeImpl`/`sameType` + # are buggy with generics + # due to https://github.com/nim-lang/Nim/issues/14021 + let lateBind_masked_fill = ident"masked_fill" + let lateBind_masked_axis_fill = ident"masked_axis_fill" + let lateBind_index_fill = ident"index_fill" + + result = quote do: + type FancyType = typeof(`selector`) + when FancyType is (array or seq): + type FancyTensorType = typeof(toTensor(`selector`)) + else: + type FancyTensorType = FancyType + when FancyTensorType is Tensor[bool]: + when FancySelectorKind(`fancy`) == FancyUnknownFull: + `lateBind_masked_fill`(`t`, `selector`, `val`) + elif FancySelectorKind(`fancy`) == FancyUnknownAxis: + `lateBind_masked_axis_fill`(`t`, `selector`, `axis`, `val`) + else: + {.error: "Unreachable".} + else: + `lateBind_index_fill`(`t`, `axis`, `selector`, `val`) # ############################################################################ # Slicing a var returns a var (for Result[_] += support) # And apply2(result[_], foo) support +# +# Unused: Nim support for var return types is problematic proc slicer_var[T](t: var AnyTensor[T], slices: varargs[SteppedSlice]): var AnyTensor[T] {.noInit,noSideEffect.}= ## Take a Tensor and SteppedSlices diff --git a/src/tensor/selectors.nim b/src/tensor/selectors.nim index b68f540eb..76b028c1e 100644 --- a/src/tensor/selectors.nim +++ b/src/tensor/selectors.nim @@ -23,7 +23,10 @@ import ./backend/metadataArray, ./higher_order_foldreduce, std/sequtils -func index_select*[T; Idx: byte or char or SomeNumber](t: Tensor[T], axis: int, indices: Tensor[Idx]): Tensor[T] {.noInit.} = +# Indexed axis +# -------------------------------------------------------------------------------------------- + +func index_select*[T; Idx: byte or char or SomeInteger](t: Tensor[T], axis: int, indices: Tensor[Idx]): Tensor[T] {.noInit.} = ## Take elements from a tensor along an axis using the indices Tensor. ## This is equivalent to NumPy `take`. ## The result does not share the input storage, there are copies. @@ -42,8 +45,44 @@ func index_select*[T; Idx: byte or char or SomeNumber](t: Tensor[T], axis: int, var t_slice = t.atAxisIndex(axis, int(index)) r_slice.copyFrom(t_slice) +func index_select*[T; Idx: byte or char or SomeInteger](t: Tensor[T], axis: int, indices: openarray[Idx]): Tensor[T] {.noInit.} = + ## Take elements from a tensor along an axis using the indices Tensor. + ## This is equivalent to NumPy `take`. + ## The result does not share the input storage, there are copies. + ## The tensors containing the indices can be an integer, byte or char tensor. + + var select_shape = t.shape + select_shape[axis] = indices.len + result = newTensorUninit[T](select_shape) + + # TODO: optim for contiguous tensors + # TODO: use OpenMP for tensors of non-ref/strings/seqs + for i, index in indices: + var r_slice = result.atAxisIndex(axis, i) + var t_slice = t.atAxisIndex(axis, int(index)) + r_slice.copyFrom(t_slice) + +proc index_fill*[T; Idx: byte or char or SomeInteger](t: var Tensor[T], axis: int, indices: Tensor[Idx], value: T) = + ## Replace elements of `t` indicated by their `indices` along `axis` with `value` + ## This is equivalent to Numpy `put`. + for i, index in enumerate(indices): + var t_slice = t.atAxisIndex(axis, int(index)) + for old_val in t_slice.mitems(): + old_val = value + +proc index_fill*[T; Idx: byte or char or SomeInteger](t: var Tensor[T], axis: int, indices: openarray[Idx], value: T) = + ## Replace elements of `t` indicated by their `indices` along `axis` with `value` + ## This is equivalent to Numpy `put`. + for i, index in indices: + var t_slice = t.atAxisIndex(axis, int(index)) + for old_val in t_slice.mitems(): + old_val = value + +# Mask full tensor +# -------------------------------------------------------------------------------------------- + func masked_select*[T](t: Tensor[T], mask: Tensor[bool]): Tensor[T] {.noInit.} = - ## Take elements from a tensor accordinf to the provided boolean mask + ## Take elements from a tensor according to the provided boolean mask ## ## Returns a **flattened** tensor which is the concatenation of values for which the mask is true. ## @@ -56,6 +95,7 @@ func masked_select*[T](t: Tensor[T], mask: Tensor[bool]): Tensor[T] {.noInit.} = size += int(val) result = newTensorUninit[T](size) + withMemoryOptimHints() var idx = 0 let dst{.restrict.} = result.dataArray @@ -65,18 +105,80 @@ func masked_select*[T](t: Tensor[T], mask: Tensor[bool]): Tensor[T] {.noInit.} = inc idx assert idx == size -func masked_axis_select*[T](t: Tensor[T], mask: Tensor[bool], axis: int): Tensor[T] {.noInit.} = - ## Take elements from a tensor according to the provided boolean mask. - ## The mask must be a 1D tensor and is applied along an axis, by default 0. +func masked_select*[T](t: Tensor[T], mask: openarray): Tensor[T] {.noInit.} = + ## Take elements from a tensor according to the provided boolean mask ## - ## The result will be the concatenation of values for which the mask is true. + ## The boolean mask must be + ## - an array or sequence of bools + ## - an array of arrays of bools, + ## - ... ## - ## For example, for a 1D tensor `t` - ## t.masked_select(t > 0) will return a tensor with - ## only the positive values of t. + ## Returns a **flattened** tensor which is the concatenation of values for which the mask is true. ## ## The result does not share input storage. - doAssert mask.shape.len == 1, "Mask must be a 1d tensor" + t.masked_select mask.toTensor() + +func masked_fill*[T](t: var Tensor[T], mask: Tensor[bool], value: T) = + ## For the index of each element of t. + ## Fill the elements at ``t[index]`` with the ``value`` + ## if their corresponding ``mask[index]`` is true. + ## If not they are untouched. + ## + ## Example: + ## + ## t.masked_fill(t > 0, -1) + ## + ## or alternatively: + ## + ## t.masked_fill(t > 0): -1 + check_elementwise(t, mask) + + # Due to requiring unnecessary assigning `x` for a `false` mask + # apply2_inline is a bit slower for very sparse mask. + # As this is a critical operation, especially on dataframes, we use the lower level construct. + # + # t.apply2_inline(mask): + # if y: + # value + # else: + # x + omp_parallel_blocks(block_offset, block_size, t.size): + for tElem, maskElem in mzip(t, mask, block_offset, block_size): + if maskElem: + tElem = value + + +func masked_fill*[T](t: var Tensor[T], mask: openarray, value: T) = + ## For the index of each element of t. + ## Fill the elements at ``t[index]`` with the ``value`` + ## if their corresponding ``mask[index]`` is true. + ## If not they are untouched. + ## + ## Example: + ## + ## t.masked_fill(t > 0, -1) + ## + ## or alternatively: + ## + ## t.masked_fill(t > 0): -1 + ## + ## The boolean mask must be + ## - an array or sequence of bools + ## - an array of arrays of bools, + ## - ... + ## + t.masked_fill(mask.toTensor(), value) + +# Mask axis +# -------------------------------------------------------------------------------------------- + +template masked_axis_select_impl[T](result: var Tensor[T], t: Tensor[T], mask: Tensor[bool] or openArray[bool], axis: int) = + ## Indirection because Nim proc can't type match "Tensor[bool] or openArray[bool]" with an array[N, bool] + when mask is Tensor: + doAssert mask.shape.len == 1, "Mask must be a 1d tensor" + doAssert t.shape[axis] == mask.shape[0], "The mask length doesn't match the axis length." + else: + doAssert t.shape[axis] == mask.len, "The mask length doesn't match the axis length." # TODO: fold_inline should accept an accumType like fold_axis_inline var size = 0 @@ -88,7 +190,7 @@ func masked_axis_select*[T](t: Tensor[T], mask: Tensor[bool], axis: int): Tensor result = newTensorUninit[T](shape) # Track the current slice of the result tensor - var dstSlice = shape.mapIt((0.. 0) will return a tensor with + ## only the positive values of t. + ## + ## The result does not share input storage. + let mask = mask.squeeze() # make 1D if coming from unreduced axis aggregation like sum + masked_axis_select_impl(result, t, mask, axis) -func masked_axis_fill*[T](t: var Tensor[T], mask: Tensor[bool], axis: int, value: T or Tensor[T]) = - ## Take a 1D boolean mask tensor with size equal to the `t.shape[axis]` - ## The axis index that are set to true in the mask will be filled with `value` +func masked_axis_select*[T](t: Tensor[T], mask: openArray[bool], axis: int): Tensor[T] {.noInit.} = + ## Take elements from a tensor according to the provided boolean mask. + ## The mask must be a 1D tensor and is applied along an axis, by default 0. + ## + ## The result will be the concatenation of values for which the mask is true. + ## + ## For example, for a 1D tensor `t` + ## t.masked_select(t > 0) will return a tensor with + ## only the positive values of t. + ## + ## The result does not share input storage. + masked_axis_select_impl(result, t, mask, axis) +template masked_axis_fill_impl[T](t: var Tensor[T], mask: Tensor[bool] or openArray[bool], axis: int, value: T or Tensor[T]) = + ## Indirection because Nim proc can't type match "Tensor[bool] or openArray[bool]" with an array[N, bool] # TODO: proper check - doAssert mask.shape.len == 1, "Mask must be a 1d tensor" + when mask is Tensor: + doAssert mask.shape.len == 1, "Mask must be a 1d tensor" + doAssert t.shape[axis] == mask.shape[0], "The mask length doesn't match the axis length." + else: + doAssert t.shape[axis] == mask.len, "The mask length doesn't match the axis length." # N-D tensor case, we iterate on t axis # We update the slice of t if mask is true. # Track the current slice of the result tensor - var dstSlice = t.shape.mapIt((0.. 0, -1) + ## Limitation: + ## If value is a Tensor, only filling via broadcastable tensors is supported at the moment + ## for example if filling axis of a tensor `t` of shape [4, 3] the corresponding shapes are valid + ## [4, 3].masked_axis_fill(mask = [1, 3], axis = 1, value = [4, 1]) ## - ## or alternatively: + ## with values + ## t = [[ 4, 99, 2], + ## [ 3, 4, 99], + ## [ 1, 8, 7], + ## [ 8, 6, 8]].toTensor() + ## mask = [false, true, true] + ## value = [[10], + ## [20], + ## [30], + ## [40]].toTensor() ## - ## t.masked_fill(t > 0): -1 - check_elementwise(t, mask) + ## result = [[ 4, 10, 10], + ## [ 3, 20, 20], + ## [ 1, 30, 30], + ## [ 8, 40, 40]].toTensor() + # TODO: support filling with a multidimensional tensor + let mask = mask.squeeze() # make 1D if coming from unreduced axis aggregation like sum + # TODO: squeeze exactly depending on axis to prevent accepting invalid values + masked_axis_fill_impl(t, mask, axis, value) - # Due to requiring unnecessary assigning `x` for a `false` mask - # apply2_inline is a bit slower for very sparse mask. - # As this is a critical operation, especially on dataframes, we use the lower level construct. - # - # t.apply2_inline(mask): - # if y: - # value - # else: - # x - omp_parallel_blocks(block_offset, block_size, t.size): - for tElem, maskElem in mzip(t, mask, block_offset, block_size): - if maskElem: - tElem = value +func masked_axis_fill*[T](t: var Tensor[T], mask: openArray[bool], axis: int, value: T or Tensor[T]) = + ## Take a 1D boolean mask tensor with size equal to the `t.shape[axis]` + ## The axis index that are set to true in the mask will be filled with `value` + ## + ## Limitation: + ## If value is a Tensor, only filling via broadcastable tensors is supported at the moment + ## for example if filling axis of a tensor `t` of shape [4, 3] the corresponding shapes are valid + ## [4, 3].masked_axis_fill(mask = [1, 3], axis = 1, value = [4, 1]) + ## + ## with values + ## t = [[ 4, 99, 2], + ## [ 3, 4, 99], + ## [ 1, 8, 7], + ## [ 8, 6, 8]].toTensor() + ## mask = [false, true, true] + ## value = [[10], + ## [20], + ## [30], + ## [40]].toTensor() + ## + ## result = [[ 4, 10, 10], + ## [ 3, 20, 20], + ## [ 1, 30, 30], + ## [ 8, 40, 40]].toTensor() + # TODO: support filling with a multidimensional tensor + masked_axis_fill_impl(t, mask, axis, value) +# Apply N-D mask along an axis +# -------------------------------------------------------------------------------------------- func masked_fill_along_axis*[T](t: var Tensor[T], mask: Tensor[bool], axis: int, value: T) = ## Take a boolean mask tensor and diff --git a/tests/_split_tests/tests_tensor_part01.nim b/tests/_split_tests/tests_tensor_part01.nim index 56c23e864..ce9f93f57 100644 --- a/tests/_split_tests/tests_tensor_part01.nim +++ b/tests/_split_tests/tests_tensor_part01.nim @@ -20,4 +20,6 @@ import ../tensor/test_operators_comparison, ../tensor/test_accessors, ../tensor/test_accessors_slicer, + ../tensor/test_selectors, + ../tensor/test_fancy_indexing, ../tensor/test_display diff --git a/tests/manual_checks/fancy_indexing.py b/tests/manual_checks/fancy_indexing.py new file mode 100644 index 000000000..8a109aade --- /dev/null +++ b/tests/manual_checks/fancy_indexing.py @@ -0,0 +1,190 @@ +import numpy as np + +def index_select(): + print('Index select') + print('--------------------------') + x = np.array([[ 4, 99, 2], + [ 3, 4, 99], + [ 1, 8, 7], + [ 8, 6, 8]]) + + print(x) + print('--------------------------') + print('x[:, [0, 2]]') + print(x[:, [0, 2]]) + print('--------------------------') + print('x[[1, 3], :]') + print(x[[1, 3], :]) + +def masked_select(): + print('Masked select') + print('--------------------------') + x = np.array([[ 4, 99, 2], + [ 3, 4, 99], + [ 1, 8, 7], + [ 8, 6, 8]]) + + print(x) + print('--------------------------') + print('x[x > 50]') + print(x[x > 50]) + print('--------------------------') + print('x[x < 50]') + print(x[x < 50]) + +def masked_axis_select(): + print('Masked axis select') + print('--------------------------') + x = np.array([[ 4, 99, 2], + [ 3, 4, 99], + [ 1, 8, 7], + [ 8, 6, 8]]) + + print(x) + print('--------------------------') + print('x[:, np.sum(x, axis = 0) > 50]') + print(x[:, np.sum(x, axis = 0) > 50]) + print('--------------------------') + print('x[np.sum(x, axis = 1) > 50, :]') + print(x[np.sum(x, axis = 1) > 50, :]) + +# index_select() +# masked_select() +# masked_axis_select() + +print('\n#########################################\n') +print('Fancy mutation') + +def index_fill(): + print('Index fill') + print('--------------------------') + x = np.array([[ 4, 99, 2], + [ 3, 4, 99], + [ 1, 8, 7], + [ 8, 6, 8]]) + + print(x) + print('--------------------------') + y = x.copy() + print('y[:, [0, 2]] = -100') + y[:, [0, 2]] = -100 + print(y) + print('--------------------------') + y = x.copy() + print('y[[1, 3], :] = -100') + y[[1, 3], :] = -100 + print(y) + print('--------------------------') + +def masked_fill(): + print('Masked fill') + print('--------------------------') + x = np.array([[ 4, 99, 2], + [ 3, 4, 99], + [ 1, 8, 7], + [ 8, 6, 8]]) + + print(x) + print('--------------------------') + y = x.copy() + print('y[y > 50] = -100') + y[y > 50] = -100 + print(y) + print('--------------------------') + y = x.copy() + print('y[y < 50] = -100') + y[y < 50] = -100 + print(y) + print('--------------------------') + +def masked_axis_fill_value(): + print('Masked axis fill with value') + print('--------------------------') + x = np.array([[ 4, 99, 2], + [ 3, 4, 99], + [ 1, 8, 7], + [ 8, 6, 8]]) + + print(x) + print('--------------------------') + y = x.copy() + print('y[:, y.sum(axis = 0) > 50] = -100') + y[:, y.sum(axis = 0) > 50] = -100 + print(y) + print('--------------------------') + y = x.copy() + print('y[y.sum(axis = 1) > 50, :] = -100') + y[y.sum(axis = 1) > 50, :] = -100 + print(y) + print('--------------------------') + +def masked_axis_fill_tensor_invalid_1(): + # ValueError: shape mismatch: + # value array of shape (4,) could not be broadcast + # to indexing result of shape (2,4) + print('Masked axis fill with tensor - invalid numpy syntax') + print('--------------------------') + x = np.array([[ 4, 99, 2], + [ 3, 4, 99], + [ 1, 8, 7], + [ 8, 6, 8]]) + + print(x) + print('--------------------------') + y = x.copy() + print('y[:, y.sum(axis = 0) > 50] = np.array([10, 20, 30, 40])') + y[:, y.sum(axis = 0) > 50] = np.array([10, 20, 30, 40]) + print(y) + +def masked_axis_fill_tensor_valid_1(): + print('Masked axis fill with tensor - 1d tensor broadcasting') + print('--------------------------') + x = np.array([[ 4, 99, 2], + [ 3, 4, 99], + [ 1, 8, 7], + [ 8, 6, 8]]) + + print(x) + print('--------------------------') + y = x.copy() + print('y[:, y.sum(axis = 0) > 50] = np.array([[10], [20], [30], [40]])') + y[:, y.sum(axis = 0) > 50] = np.array([[10], [20], [30], [40]]) + print(y) + print('--------------------------') + y = x.copy() + print('y[y.sum(axis = 1) > 50, :] = np.array([-10, -20, -30])') + y[y.sum(axis = 1) > 50, :] = np.array([-10, -20, -30]) + print(y) + print('--------------------------') + +def masked_axis_fill_tensor_valid_2(): + print('Masked axis fill with tensor - multidimensional tensor') + print('--------------------------') + x = np.array([[ 4, 99, 2], + [ 3, 4, 99], + [ 1, 8, 7], + [ 8, 6, 8]]) + + print(x) + print('--------------------------') + y = x.copy() + print('y[:, y.sum(axis = 0) > 50] = np.array([[10, 50], [20, 60], [30, 70], [40, 80]])') + y[:, y.sum(axis = 0) > 50] = np.array([[10, 50], + [20, 60], + [30, 70], + [40, 80]]) + print(y) + print('--------------------------') + y = x.copy() + print('y[y.sum(axis = 1) > 50, :] = np.array([-10, -20, -30], [-40, -50, -60])') + y[y.sum(axis = 1) > 50, :] = np.array([[-10, -20, -30], + [-40, -50, -60]]) + print(y) + print('--------------------------') + +# index_fill() +# masked_fill() +# masked_axis_fill_value() +masked_axis_fill_tensor_invalid_1() +# masked_axis_fill_tensor_valid_1() +# masked_axis_fill_tensor_valid_2() diff --git a/tests/tensor/test_fancy_indexing.nim b/tests/tensor/test_fancy_indexing.nim new file mode 100644 index 000000000..f0dc28841 --- /dev/null +++ b/tests/tensor/test_fancy_indexing.nim @@ -0,0 +1,203 @@ +# Copyright 2017-2020 Mamy André-Ratsimbazafy +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import ../../src/arraymancer +import unittest + +suite "Fancy indexing": + # x = np.array([[ 4, 99, 2], + # [ 3, 4, 99], + # [ 1, 8, 7], + # [ 8, 6, 8]]) + + let x = [[ 4, 99, 2], + [ 3, 4, 99], + [ 1, 8, 7], + [ 8, 6, 8]].toTensor() + + test "Index selection via fancy indexing": + block: # print(x[:, [0, 2]]) + let r = x[_, [0, 2]] + + let exp = [[ 4, 2], + [ 3, 99], + [ 1, 7], + [ 8, 8]].toTensor() + + check: r == exp + + block: # print(x[[1, 3], :]) + let r = x[[1, 3], _] + + let exp = [[3, 4, 99], + [8, 6, 8]].toTensor() + + check: r == exp + + test "Masked selection via fancy indexing": + block: + let r = x[x >. 50] + let exp = [99, 99].toTensor() + check: r == exp + + block: + let r = x[x <. 50] + let exp = [4, 2, 3, 4, 1, 8, 7, 8, 6, 8].toTensor() + check: r == exp + + test "Masked axis selection via fancy indexing": + block: # print('x[:, np.sum(x, axis = 0) > 50]') + let r = x[_, x.sum(axis = 0) >. 50] + + let exp = [[99, 2], + [ 4, 99], + [ 8, 7], + [ 6, 8]].toTensor() + + check: r == exp + + block: # print('x[np.sum(x, axis = 1) > 50, :]') + let r = x[x.sum(axis = 1) >. 50, _] + + let exp = [[4, 99, 2], + [3, 4, 99]].toTensor() + + check: r == exp + + test "Index assign value via fancy indexing": + block: # y[:, [0, 2]] = -100 + var y = x.clone() + y[_, [0, 2]] = -100 + + let exp = [[-100, 99, -100], + [-100, 4, -100], + [-100, 8, -100], + [-100, 6, -100]].toTensor() + + check: y == exp + + block: # y[[1, 3], :] = -100 + var y = x.clone() + y[[1, 3], _] = -100 + + let exp = [[ 4, 99, 2], + [-100, -100, -100], + [ 1, 8, 7], + [-100, -100, -100]].toTensor() + + check: y == exp + + test "Masked assign value via fancy indexing": + block: # y[y > 50] = -100 + var y = x.clone() + y[y >. 50] = -100 + + let exp = [[ 4, -100, 2], + [ 3, 4, -100], + [ 1, 8, 7], + [ 8, 6, 8]].toTensor() + + check: y == exp + + block: # y[y < 50] = -100 + var y = x.clone() + y[y <. 50] = -100 + + let exp = [[ -100, 99, -100], + [ -100, -100, 99], + [ -100, -100, -100], + [ -100, -100, -100]].toTensor() + + check: y == exp + + test "Masked axis assign value via fancy indexing": + block: # y[:, y.sum(axis = 0) > 50] = -100 + var y = x.clone() + y[_, y.sum(axis = 0) >. 50] = -100 + + let exp = [[ 4, -100, -100], + [ 3, -100, -100], + [ 1, -100, -100], + [ 8, -100, -100]].toTensor() + + check: y == exp + + block: # y[y.sum(axis = 1) > 50, :] = -100 + var y = x.clone() + y[y.sum(axis = 1) >. 50, _] = -100 + + let exp = [[-100, -100, -100], + [-100, -100, -100], + [ 1, 8, 7], + [ 8, 6, 8]].toTensor() + + check: y == exp + + test "Masked axis assign tensor via fancy indexing - invalid Numpy syntaxes": + block: # y[:, y.sum(axis = 0) > 50] = np.array([10, 20, 30, 40]) + var y = x.clone() + + expect(IndexError): + y[_, y.sum(axis = 0) >. 50] = [10, 20, 30, 40].toTensor() + + test "Masked axis assign broadcastable 1d tensor via fancy indexing": + block: # y[:, y.sum(axis = 0) > 50] = np.array([[10], [20], [30], [40]]) + var y = x.clone() + y[_, y.sum(axis = 0) >. 50] = [[10], [20], [30], [40]].toTensor() + + let exp = [[ 4, 10, 10], + [ 3, 20, 20], + [ 1, 30, 30], + [ 8, 40, 40]].toTensor() + + check: y == exp + + block: # y[y.sum(axis = 1) > 50, :] = np.array([-10, -20, -30]) + var y = x.clone() + y[y.sum(axis = 1) >. 50, _] = [[-10, -20, -30]].toTensor() + + let exp = [[-10, -20, -30], + [-10, -20, -30], + [ 1, 8, 7], + [ 8, 6, 8]].toTensor() + + check: y == exp + + # TODO - only broadcastable tensor assign are supported at the moment + # test "Masked axis assign multidimensional tensor via fancy indexing": + # block: # y[:, y.sum(axis = 0) > 50] = np.array([[10, 50], [20, 60], [30, 70], [40, 80]]) + # var y = x.clone() + # y[_, y.sum(axis = 0) >. 50] = [[10, 50], + # [20, 60], + # [30, 70], + # [40, 80]].toTensor() + # + # let exp = [[ 4, 10, 50], + # [ 3, 20, 60], + # [ 1, 30, 70], + # [ 8, 40, 80]].toTensor() + # + # check: y == exp + # + # block: # y[y.sum(axis = 1) > 50, :] = np.array([-10, -20, -30], [-40, -50, -60]) + # var y = x.clone() + # y[y.sum(axis = 1) >. 50, _] = [[-10, -20, -30], + # [-40, -50, -60]].toTensor() + # + # let exp = [[-10, -20, -30], + # [-40, -50, -60], + # [ 1, 8, 7], + # [ 8, 6, 8]].toTensor() + # + # check: y == exp diff --git a/tests/tensor/test_selectors.nim b/tests/tensor/test_selectors.nim index 24fabc005..6710a240d 100644 --- a/tests/tensor/test_selectors.nim +++ b/tests/tensor/test_selectors.nim @@ -40,6 +40,98 @@ suite "Selectors": x.index_select(axis = 0, indices) == ax0 x.index_select(axis = 1, indices) == ax1 + # ------------------------------------------------------ + # Selection with regular arrays/sequences + + block: # Numpy + let a = [4, 3, 5, 7, 6, 8].toTensor + let indices = [0, 1, 4] + + check: a.index_select(axis = 0, indices) == [4, 3, 6].toTensor + + block: # PyTorch + let x = [[ 0.1427, 0.0231, -0.5414, -1.0009], + [-0.4664, 0.2647, -0.1228, -1.1068], + [-1.1734, -0.6571, 0.7230, -0.6004]].toTensor + + let indices = [0, 2] + + let ax0 = [[ 0.1427, 0.0231, -0.5414, -1.0009], + [-1.1734, -0.6571, 0.7230, -0.6004]].toTensor + let ax1 = [[ 0.1427, -0.5414], + [-0.4664, -0.1228], + [-1.1734, 0.7230]].toTensor + + check: + x.index_select(axis = 0, indices) == ax0 + x.index_select(axis = 1, indices) == ax1 + + test "Index_fill (Numpy put)": + block: # Numpy + var a = [4, 3, 5, 7, 6, 8].toTensor + let indices = [0, 1, 4].toTensor + + a.index_fill(axis = 0, indices, -1) + check: a == [-1, -1, 5, 7, -1, 8].toTensor + + block: # PyTorch + let x = [[ 0.1427, 0.0231, -0.5414, -1.0009], + [-0.4664, 0.2647, -0.1228, -1.1068], + [-1.1734, -0.6571, 0.7230, -0.6004]].toTensor + + let indices = [0, 2].toTensor + + var x0 = x.clone() + var x1 = x.clone() + + x0.index_fill(axis = 0, indices, -10.0) + x1.index_fill(axis = 1, indices, -10.0) + + let ax0 = [[ -10.0 , -10.0 , -10.0 , -10.0 ], + [ -0.4664, 0.2647, -0.1228, -1.1068], + [ -10.0 , -10.0 , -10.0 , -10.0 ]].toTensor + let ax1 = [[-10.0, 0.0231, -10.0, -1.0009], + [-10.0, 0.2647, -10.0, -1.1068], + [-10.0, -0.6571, -10.0, -0.6004]].toTensor + + check: + x0 == ax0 + x1 == ax1 + + # ------------------------------------------------------ + # Selection with regular arrays/sequences + + block: # Numpy + var a = [4, 3, 5, 7, 6, 8].toTensor + let indices = [0, 1, 4] + + a.index_fill(axis = 0, indices, -1) + check: a == [-1, -1, 5, 7, -1, 8].toTensor + + block: # PyTorch + let x = [[ 0.1427, 0.0231, -0.5414, -1.0009], + [-0.4664, 0.2647, -0.1228, -1.1068], + [-1.1734, -0.6571, 0.7230, -0.6004]].toTensor + + let indices = [0, 2] + + var x0 = x.clone() + var x1 = x.clone() + + x0.index_fill(axis = 0, indices, -10.0) + x1.index_fill(axis = 1, indices, -10.0) + + let ax0 = [[ -10.0 , -10.0 , -10.0 , -10.0 ], + [ -0.4664, 0.2647, -0.1228, -1.1068], + [ -10.0 , -10.0 , -10.0 , -10.0 ]].toTensor + let ax1 = [[-10.0, 0.0231, -10.0, -1.0009], + [-10.0, 0.2647, -10.0, -1.1068], + [-10.0, -0.6571, -10.0, -0.6004]].toTensor + + check: + x0 == ax0 + x1 == ax1 + test "Masked_select": block: # Numpy reference doc # https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#boolean-array-indexing @@ -56,6 +148,20 @@ suite "Selectors": let expected = [1.0, 2.0, 3.0].toTensor() check: r == expected + block: # with regular arrays/sequences + let x = [[1.0, 2.0], + [Nan, 3.0], + [Nan, Nan]].toTensor + + let r = x.masked_select( + [[true, true], + [false, true], + [false, false]] + ) + + let expected = [1.0, 2.0, 3.0].toTensor() + check: r == expected + test "Masked_fill": block: # Numpy reference doc # https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#boolean-array-indexing @@ -73,6 +179,21 @@ suite "Selectors": let expected = [[1.0, 2.0], [-1.0, 3.0], [-1.0, -1.0]].toTensor() check: x == expected + block: # with regular arrays/sequences + var x = [[1.0, 2.0], + [Nan, 3.0], + [Nan, Nan]].toTensor + + x.masked_fill( + [[false, false], + [true, false], + [true, true]], + -1.0 + ) + + let expected = [[1.0, 2.0], [-1.0, 3.0], [-1.0, -1.0]].toTensor() + check: x == expected + test "Masked_axis_select": block: # Numpy reference doc # https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#boolean-array-indexing @@ -88,7 +209,7 @@ suite "Selectors": [2, 2]].toTensor let rowsum = x.sum(axis = 1) - let cond = rowsum .<= 2 + let cond = rowsum <=. 2 let r = x.masked_axis_select(cond.squeeze(), axis = 0) let expected = [[0, 1], @@ -96,7 +217,22 @@ suite "Selectors": check: r == expected - test "masked_axis_fill": + block: # With regular arrays/sequences + + let x = [[0, 1], + [1, 1], + [2, 2]].toTensor + + let rowsum = x.sum(axis = 1) + let cond = [true, + true, + false] + let r = x.masked_axis_select(cond, axis = 0) + + let expected = [[0, 1], + [1, 1]].toTensor + + test "Masked_axis_fill with value": block: # Numpy # Fill all columns which sum up to greater than 1 # with -10 @@ -109,7 +245,7 @@ suite "Selectors": [ 1, 2, 0], [ 1, -1, 1]].toTensor - let cond = squeeze(a.sum(axis = 0) .> 1) + let cond = squeeze(a.sum(axis = 0) >. 1) a.masked_axis_fill(cond, axis = 1, -10) let expected = [[-1, -2, -10], @@ -118,7 +254,22 @@ suite "Selectors": check: a == expected - block: # Fill with tensor + block: # With regular arrays/sequences + var a = [[-1, -2, 1], + [ 1, 2, 0], + [ 1, -1, 1]].toTensor + + let cond = [false, false, true] + a.masked_axis_fill(cond, axis = 1, -10) + + let expected = [[-1, -2, -10], + [ 1, 2, -10], + [ 1, -1, -10]].toTensor + + check: a == expected + + test "Masked_axis_fill with tensor": + block: # import numpy as np # a = np.array([[-1, -2, 1], [1, 2, 0], [1, -1, 1]]) # print(a.sum(axis=0) > 1) @@ -131,7 +282,23 @@ suite "Selectors": let b = [-10, -20, -30].toTensor.unsqueeze(1) - let cond = squeeze(a.sum(axis = 0) .> 1) + let cond = squeeze(a.sum(axis = 0) >. 1) + a.masked_axis_fill(cond, axis = 1, b) + + let expected = [[-1, -2, -10], + [ 1, 2, -20], + [ 1, -1, -30]].toTensor + + check: a == expected + + block: # With regular arrays/sequences + var a = [[-1, -2, 1], + [ 1, 2, 0], + [ 1, -1, 1]].toTensor + + let b = [-10, -20, -30].toTensor.unsqueeze(1) + + let cond = [false, false, true] a.masked_axis_fill(cond, axis = 1, b) let expected = [[-1, -2, -10], @@ -146,7 +313,7 @@ suite "Selectors": [ 1, 2, 0], [ 1, -1, 1]].toTensor - a.masked_fill_along_axis(a.sum(axis = 0) .> 1, axis = 0, -10) + a.masked_fill_along_axis(a.sum(axis = 0) >. 1, axis = 0, -10) let expected = [[-1, -2, -10], [ 1, 2, -10], diff --git a/tests/tests_cpu.nim b/tests/tests_cpu.nim index a7beb426e..f343fda05 100644 --- a/tests/tests_cpu.nim +++ b/tests/tests_cpu.nim @@ -18,6 +18,8 @@ import ../src/arraymancer, ./tensor/test_operators_comparison, ./tensor/test_accessors, ./tensor/test_accessors_slicer, + ./tensor/test_selectors, + ./tensor/test_fancy_indexing, ./tensor/test_display, ./tensor/test_operators_blas, ./tensor/test_math_functions,