Search in sources :

Example 1 with IntKeyVector

use of com.tencent.angel.ml.math2.vector.IntKeyVector in project angel by Tencent.

the class UpdateColsParam method split.

@Override
public List<PartitionUpdateParam> split() {
    List<PartitionKey> pkeys = PSAgentContext.get().getMatrixMetaManager().getPartitions(matrixId);
    List<PartitionUpdateParam> params = new ArrayList<>();
    int start = 0, end = 0;
    for (PartitionKey pkey : pkeys) {
        long startCol = pkey.getStartCol();
        long endCol = pkey.getEndCol();
        if (start < ((IntKeyVector) cols).getDim() && VectorUtils.getLong(cols, start) >= startCol) {
            while (end < ((IntKeyVector) cols).getDim() && VectorUtils.getLong(cols, end) < endCol) end++;
            long[] part = new long[end - start];
            if (cols instanceof IntIntVector) {
                ArrayCopy.copy(((IntIntVector) cols).getStorage().getValues(), start, part, 0, end - start);
            } else {
                System.arraycopy(((IntLongVector) cols).getStorage().getValues(), start, part, 0, end - start);
            }
            long firstKey = 0l;
            for (Map.Entry<Long, Vector> first : values.entrySet()) {
                firstKey = first.getKey();
                break;
            }
            if (values.get(firstKey) instanceof IntDoubleVector) {
                IntDoubleVector[] updates = new IntDoubleVector[part.length];
                for (int i = 0; i < part.length; i++) updates[i] = (IntDoubleVector) values.get(part[i]);
                params.add(new PartitionUpdateColsParam(matrixId, pkey, rows, part, VFactory.compIntDoubleVector(rows.length, updates, part.length), op));
            } else if (values.get(firstKey) instanceof IntFloatVector) {
                IntFloatVector[] updates = new IntFloatVector[part.length];
                for (int i = 0; i < part.length; i++) updates[i] = (IntFloatVector) values.get(part[i]);
                params.add(new PartitionUpdateColsParam(matrixId, pkey, rows, part, VFactory.compIntFloatVector(rows.length, updates, part.length), op));
            } else {
                throw new AngelException("Update data type should be float or double!");
            }
            start = end;
        }
    }
    return params;
}
Also used : AngelException(com.tencent.angel.exception.AngelException) PartitionUpdateParam(com.tencent.angel.ml.matrix.psf.update.base.PartitionUpdateParam) PartitionKey(com.tencent.angel.PartitionKey) Vector(com.tencent.angel.ml.math2.vector.Vector)

Aggregations

PartitionKey (com.tencent.angel.PartitionKey)1 AngelException (com.tencent.angel.exception.AngelException)1 Vector (com.tencent.angel.ml.math2.vector.Vector)1 PartitionUpdateParam (com.tencent.angel.ml.matrix.psf.update.base.PartitionUpdateParam)1