def transform_PrimCall(self, expr):
    args = self.transform_args(expr.args)

    arg_types = get_types(args)
    if all(isinstance(t, ScalarT) for t in arg_types):
      upcast_types = expr.prim.expected_input_types(arg_types)
      result_type = expr.prim.result_type(upcast_types)
      return syntax.PrimCall(expr.prim, args, type = result_type)
    elif expr.prim == prims.is_:
      if arg_types[0] != arg_types[1]:
        return false
      elif arg_types[0] == NoneType:
        return true
      else:
        return syntax.PrimCall(prims.is_, args, type = Bool)
    elif all(t.rank == 0 for t in arg_types):
      # arguments should then be tuples
      assert len(arg_types) == 2, \
        "Expected two arguments but got [%s] in %s" % (", ".join(str(t) for t in arg_types), expr)
      xt, yt = arg_types
      x, y = args
      assert xt.__class__ is TupleT, \
        "Unexpected argument types (%s,%s) for operator %s" % (xt, yt, expr.prim)
      assert yt.__class__ is TupleT, \
        "Unexepcted argument types (%s,%s) for operator %s" % (xt, yt, expr.prim)
      x_elts = self.tuple_elts(x)
      y_elts = self.tuple_elts(y)
      nx = len(x_elts)
      ny = len(y_elts)
      if expr.prim is prims.add:
        return self.tuple( tuple(x_elts) + tuple(y_elts) )
      elif expr.prim is prims.equal:
        assert len(x_elts) == len(y_elts), \
          "Can't compare tuple of unequal lengths %d and %d" % (nx, ny)
        result = true  
        for (xi, yi) in zip(x_elts, y_elts):
          elts_eq = syntax.PrimCall(prims.equal, (xi, yi), type=Bool)
          result = syntax.PrimCall(prims.logical_and, (result, elts_eq), type=Bool) 
        return result  
      elif expr.prim is prims.not_equal:
        assert len(x_elts) == len(y_elts), \
          "Can't compare tuple of unequal lengths %d and %d" % (nx, ny)
        result = false  
        for (xi, yi) in zip(x_elts, y_elts):
          elts_eq = syntax.PrimCall(prims.not_equal, (xi, yi), type=Bool)
          result = syntax.PrimCall(prims.logical_or, (result, elts_eq), type=Bool) 
        return result  
      
      else:
        assert False, "Unsupported tuple operation %s" % expr  
    else:
      assert all(t.__class__ is not NoneT for t in arg_types), \
        "Invalid argument types for prim %s: %s" % (expr.prim, arg_types,)
      elt_types = [t.elt_type if isinstance(t, ArrayT) else t for t in arg_types]
      typed_prim_fn = mk_prim_fn(expr.prim, elt_types)
      max_rank = adverb_helpers.max_rank(arg_types)
      if max_rank == 0:
        return self.call(typed_prim_fn, args)
      elif max_rank > 1:
        axis = zero_i64 
      else:
        axis = none 
      result_t = make_array_type(typed_prim_fn.return_type, max_rank)
      return Map(fn = typed_prim_fn, args = args, type = result_t, axis = axis)
dataset/ETHPy150Open iskandr/parakeet/parakeet/type_inference/type_inference.py/TypeInference.transform_PrimCall/operator `==` -> `>`