Boolean indexing with ND4J

Question

What I want to achieve:

With NumPy, we have boolean indexing which operates like this:

import numpy as np
a = np.arange(10)
mask = a < 5
a[mask]

# array([0, 1, 2, 3, 4])

What I tried:

I've been trying to make similar functionality work with nd4j:

        int[] assignments = {0,0,0,1,0,2,2};
        int[] indexes     = {0,1,2,3,4,5,7};

        INDArray asarray = Nd4j.createFromArray(assignments);
        INDArray idxarray = Nd4j.createFromArray(indexes);

        int i = 0; 
        INDArray mask = asarray.match(i, Conditions.equals());
        INDArray cluster = BooleanIndexing.applyMask(idxarray, mask);

However, this method returns cluster as a vector the same length as the mask, where all non-matching elements are set to 0. This means that any true zero values get lost. Moreover, if you set i > 0, .match() returns a boolean mask that is incorrect. The code above returns [ false, false, false, true, false, true, true] for i=1, for example.

I have found that array.eq(i) returns the correct boolean array for any i. But that still leaves the issue with the returned zeros.

Question

Does ND4J support boolean indexing? Specifically a method or chain of methods that will efficiently mimic the python code above?