def transform_IndexReduce(self, expr):
    shape = self.transform_expr(expr.shape)
    map_fn_closure = self.transform_fn(expr.fn if expr.fn else untyped_identity_function)
    combine_closure = self.transform_fn(expr.combine)
    init = self.transform_if_expr(expr.init)
    shape_t = shape.type
    if isinstance(shape_t, IntT):
      shape = self.cast(shape, Int64)
      n_indices = 1
    else:
      assert isinstance(shape_t, TupleT)
      assert all(isinstance(t, ScalarT) for t in shape_t.elt_types)
      n_indices = len(shape_t.elt_types)
      if not all(t == Int64 for t in shape_t.elt_types):
        elts = tuple(self.cast(elt, Int64) for elt in self.tuple_elts(shape))
        shape = self.tuple(elts)
    result_type, typed_fn, typed_combine = \
      specialize_IndexReduce(map_fn_closure.type, combine_closure, n_indices, init)
    if not self.is_none(init):
      init = self.cast(init, result_type)
    return syntax.IndexReduce(shape = shape, 
                              fn = make_typed_closure(map_fn_closure, typed_fn),
                              combine = make_typed_closure(combine_closure, typed_combine),
                              init = init,  
                              type = result_type)
dataset/ETHPy150Open iskandr/parakeet/parakeet/type_inference/type_inference.py/TypeInference.transform_IndexReduce/original