def specialize_AllPairs(fn, xtype, ytype): x_elt_t = array_type.lower_rank(xtype, 1) y_elt_t = array_type.lower_rank(ytype, 1) typed_map_fn = specialize(fn, [x_elt_t, y_elt_t]) elt_result_t = typed_map_fn.return_type result_t = array_type.increase_rank(elt_result_t, 2) return result_t, typed_map_fn
def specialize_Scan(map_fn, combine_fn, emit_fn, array_types, init_type = None):
acc_type, typed_map_fn, typed_combine_fn = \
specialize_Reduce(map_fn, combine_fn, array_types, init_type)
typed_emit_fn = specialize(emit_fn, [acc_type])
result_type = array_type.increase_rank(typed_emit_fn.return_type, 1)
return result_type, typed_map_fn, typed_combine_fn, typed_emit_fn
