def transform_IndexMap(self, expr):
    shape = self.transform_expr(expr.shape)
    if not isinstance(shape.type, TupleT):
      assert isinstance(shape.type, ScalarT), "Invalid shape for IndexMap: %s : %s" % (shape, shape.type)
      shape = self.tuple((shape,))
    closure = self.transform_fn(expr.fn)
    shape_t = shape.type
    if isinstance(shape_t, IntT):
      shape = self.cast(shape, Int64)
      n_indices = 1
    else:
      assert isinstance(shape_t, TupleT), "Expected shape to be tuple, instead got %s" % (shape_t,)
      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 = specialize_IndexMap(closure.type, n_indices)
    return syntax.IndexMap(shape = shape, 
                           fn = make_typed_closure(closure, typed_fn), 
                           type = result_type)
dataset/ETHPy150Open iskandr/parakeet/parakeet/type_inference/type_inference.py/TypeInference.transform_IndexMap/original