def output_type(self):
    arg = self.arg
    return rlz.shape_like(arg, dt.Array(arg.type().key_type))
def output_type(self):
    arg = self.arg
    return rlz.shape_like(arg, dt.Array(arg.type().key_type))