"""
Lift ckernels to their appropriate rank so they always consume the full array
arguments.
"""

from __future__ import absolute_import, division, print_function

from dynd import ndt, _lowlevel

from ..traversal import visit


class CKernelLifter(object):
    """
    Lift ckernels to their appropriate rank so they always consume the
    full array arguments.

    If the environment defines 'stream-outer' as True, then the
    outermost dimension is skipped, so that the operation can be
    chunked along that dimension.
    """
    def __init__(self, env):
        self.env = env

    def get_arg_type(self, arg):
        dynd_types = self.env['dynd-types']
        if arg in dynd_types:
            return dynd_types[arg]
        else:
            return ndt.type(str(arg.type))

    def op_ckernel(self, op):
        op_ndim = len(op.type.shape)
        if op.metadata['rank'] < op_ndim:
            result_ndim = self.env.get('result-ndim', 0)
            ckernel, args = op.args
            in_types = [self.get_arg_type(arg) for arg in args[1:]]
            out_type = ndt.type(str(args[0].type))

            # Replace the leading dimension type with 'strided' in each operand
            # if we're streaming it for processing BLZ
            if self.env.get('stream-outer', False) and result_ndim == op_ndim:
                # TODO: Add dynd tp.subarray(N) function like datashape has
                for i, tp in enumerate(in_types):
                    if tp.ndim == result_ndim:
                        in_types[i] = ndt.make_strided_dim(tp.element_type)
                out_type = ndt.make_strided_dim(out_type.element_type)

            op.args[0] = _lowlevel.lift_ckernel_deferred(ckernel,
                            [out_type] + in_types)


def run(func, env):
    visit(CKernelLifter(env), func)

@@ -852,54 +852,8 @@
 pe)%0A
-        if op.metadata%5B'rank'%5D %3C op_ndim:%0A    
     
@@ -909,20 +909,16 @@
         
-    
 ckernel,
@@ -933,20 +933,16 @@
 op.args%0A
-    
         
@@ -997,36 +997,32 @@
 gs%5B1:%5D%5D%0A        
-    
 out_type = ndt.t
@@ -1041,24 +1041,822 @@
 %5B0%5D.type))%0A%0A
+        if isinstance(ckernel, dict):%0A            tag = ckernel%5B'tag'%5D%0A            if tag == 'reduction':%0A                ck = ckernel%5B'ckernel'%5D%0A                assoc = ckernel%5B'assoc'%5D%0A                comm = ckernel%5B'comm'%5D%0A                ident = ckernel%5B'ident'%5D%0A                axis = ckernel%5B'axis'%5D%0A                keepdims = ckernel%5B'keepdims'%5D%0A                op.args%5B0%5D = _lowlevel.lift_reduction_ckernel_deferred(%0A                                ck, in_types%5B0%5D,%0A                                axis=axis, keepdims=keepdims,%0A                                associative=assoc, commutative=comm,%0A                                reduction_identity=ident)%0A            else:%0A                raise RuntimeError('unnrecognized ckernel tag %25s' %25 tag)%0A        elif op.metadata%5B'rank'%5D %3C op_ndim:%0A
             
@@ -2444,16 +2444,49 @@
 kernel,%0A
+                                 
         

Update execution engine to be able to execute the reduction kernels