Ejemplo n.º 1
0
    def blockmatrix_irs(self):
        scalar_ir = ir.F64(2)
        vector_ir = ir.MakeArray([ir.F64(3), ir.F64(2)],
                                 hl.tarray(hl.tfloat64))

        read = ir.BlockMatrixRead(
            ir.BlockMatrixNativeReader(resource('blockmatrix_example/0')))
        add_two_bms = ir.BlockMatrixMap2(
            read, read, 'l', 'r',
            ir.ApplyBinaryPrimOp('+', ir.Ref('l'), ir.Ref('r')), "Union")
        negate_bm = ir.BlockMatrixMap(
            read, 'element', ir.ApplyUnaryPrimOp('-', ir.Ref('element')),
            False)
        sqrt_bm = ir.BlockMatrixMap(
            read, 'element',
            hl.sqrt(construct_expr(ir.Ref('element'), hl.tfloat64))._ir, False)
        persisted = ir.BlockMatrixRead(ir.BlockMatrixPersistReader('x', read))

        scalar_to_bm = ir.ValueToBlockMatrix(scalar_ir, [1, 1], 1)
        col_vector_to_bm = ir.ValueToBlockMatrix(vector_ir, [2, 1], 1)
        row_vector_to_bm = ir.ValueToBlockMatrix(vector_ir, [1, 2], 1)
        broadcast_scalar = ir.BlockMatrixBroadcast(scalar_to_bm, [], [2, 2],
                                                   256)
        broadcast_col = ir.BlockMatrixBroadcast(col_vector_to_bm, [0], [2, 2],
                                                256)
        broadcast_row = ir.BlockMatrixBroadcast(row_vector_to_bm, [1], [2, 2],
                                                256)
        transpose = ir.BlockMatrixBroadcast(broadcast_scalar, [1, 0], [2, 2],
                                            256)
        matmul = ir.BlockMatrixDot(broadcast_scalar, transpose)

        rectangle = ir.Literal(hl.tarray(hl.tint64), [0, 1, 5, 6])
        band = ir.Literal(hl.ttuple(hl.tint64, hl.tint64), (-1, 1))
        intervals = ir.Literal(
            hl.ttuple(hl.tarray(hl.tint64), hl.tarray(hl.tint64)),
            ([0, 1, 5, 6], [5, 6, 8, 9]))

        sparsify1 = ir.BlockMatrixSparsify(read, rectangle,
                                           ir.RectangleSparsifier)
        sparsify2 = ir.BlockMatrixSparsify(read, band, ir.BandSparsifier(True))
        sparsify3 = ir.BlockMatrixSparsify(read, intervals,
                                           ir.RowIntervalSparsifier(True))

        densify = ir.BlockMatrixDensify(read)

        pow_ir = (construct_expr(ir.Ref('l'), hl.tfloat64)**construct_expr(
            ir.Ref('r'), hl.tfloat64))._ir
        squared_bm = ir.BlockMatrixMap2(scalar_to_bm, scalar_to_bm, 'l', 'r',
                                        pow_ir, "NeedsDense")
        slice_bm = ir.BlockMatrixSlice(
            matmul, [slice(0, 2, 1), slice(0, 1, 1)])

        return [
            read, persisted, add_two_bms, negate_bm, sqrt_bm, scalar_to_bm,
            col_vector_to_bm, row_vector_to_bm, broadcast_scalar,
            broadcast_col, broadcast_row, squared_bm, transpose, sparsify1,
            sparsify2, sparsify3, densify, matmul, slice_bm
        ]
Ejemplo n.º 2
0
    def blockmatrix_irs(self):
        scalar_ir = ir.F64(2)
        vector_ir = ir.MakeArray([ir.F64(3), ir.F64(2)],
                                 hl.tarray(hl.tfloat64))

        read = ir.BlockMatrixRead(
            ir.BlockMatrixNativeReader(resource('blockmatrix_example/0')))
        add_two_bms = ir.BlockMatrixMap2(
            read, read, 'l', 'r',
            ir.ApplyBinaryPrimOp('+', ir.Ref('l'), ir.Ref('r')))
        negate_bm = ir.BlockMatrixMap(
            read, 'element', ir.ApplyUnaryPrimOp('-', ir.Ref('element')))
        sqrt_bm = ir.BlockMatrixMap(
            read, 'element',
            hl.sqrt(construct_expr(ir.Ref('element'), hl.tfloat64))._ir)

        scalar_to_bm = ir.ValueToBlockMatrix(scalar_ir, [1, 1], 1)
        col_vector_to_bm = ir.ValueToBlockMatrix(vector_ir, [2, 1], 1)
        row_vector_to_bm = ir.ValueToBlockMatrix(vector_ir, [1, 2], 1)
        broadcast_scalar = ir.BlockMatrixBroadcast(scalar_to_bm, [], [2, 2],
                                                   256)
        broadcast_col = ir.BlockMatrixBroadcast(col_vector_to_bm, [0], [2, 2],
                                                256)
        broadcast_row = ir.BlockMatrixBroadcast(row_vector_to_bm, [1], [2, 2],
                                                256)
        transpose = ir.BlockMatrixBroadcast(broadcast_scalar, [1, 0], [2, 2],
                                            256)
        matmul = ir.BlockMatrixDot(broadcast_scalar, transpose)

        pow_ir = (construct_expr(ir.Ref('l'), hl.tfloat64)**construct_expr(
            ir.Ref('r'), hl.tfloat64))._ir
        squared_bm = ir.BlockMatrixMap2(scalar_to_bm, scalar_to_bm, 'l', 'r',
                                        pow_ir)
        slice_bm = ir.BlockMatrixSlice(
            matmul, [slice(0, 2, 1), slice(0, 1, 1)])

        return [
            read, add_two_bms, negate_bm, sqrt_bm, scalar_to_bm,
            col_vector_to_bm, row_vector_to_bm, broadcast_scalar,
            broadcast_col, broadcast_row, squared_bm, transpose, matmul,
            slice_bm
        ]