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