def _map_fn(x, y, z):
      return math_ops.square(x), math_ops.square(y), math_ops.square(z)