Project: 017_Residual_Deep_Gaussian_Processes_on_Manifolds | Function: sphere_expmap (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Spherical exponential map implementation for manifold-to-manifold layer computation with first-order approximation for numerical stability

**KEY TERMS:** {
  "theta": "θ = ||v|| tangent vector norm",
  "true_expmap": "exp_x exponential map formula",
  "first_order_approx": "numerical approximation for small θ",
  "jnp.cos(theta)": "cos(θ) trigonometric component",
  "jnp.sin(theta)": "sin(θ) trigonometric component"
}

**CONTEXT CODE:** 
import jax
import jax.numpy as jnp
import numpy as np
from jax.tree_util import Partial
from jax.scipy.special import gammaln


###
# GENERAL UTILS
### 


def jax_key_to_numpy_generator(key: Key) -> Generator:
    return np.random.default_rng(np.asarray(key._base_array))


def array(x):
    return jnp.array(x, dtype=jnp.float64)


@jax.jit 
def comb(N, k) :
    return jnp.round(jnp.exp(gammaln(N + 1) - gammaln(k + 1) - gammaln(N - k + 1))).astype(jnp.int64)


### 
# SPHERE UTILS
###

@jax.jit
def car_to_sph(car):
    x, y, z = car[..., 0], car[..., 1], car[..., 2]
    colat = jnp.arccos(z)
    lon = jnp.arctan2(y, x)
    lon = (lon + 2 * jnp.pi) % (2 * jnp.pi)
    return jnp.stack([colat, lon], axis=-1)


@jax.jit
def sph_to_car(sph) :
    colat, lon = sph[..., 0], sph[..., 1]
    z = jnp.cos(colat)
    r = jnp.sin(colat)
    x = r * jnp.cos(lon)
    y = r * jnp.sin(lon)
    return jnp.stack([x, y, z], axis=-1)


@jax.jit
def sphere_expmap(x, v) :
    raise NotImplementedError("Function implementation has been masked for testing")

@jax.jit 
def sphere_to_tangent(x, v) :
    v_x = jnp.sum(x * v, axis=-1, keepdims=True)
    return v - v_x * x


@Partial(jax.jit, static_argnames=("sphere_dim"))
def num_phases_in_frequency(sphere_dim: int, frequency) :
    l, d = frequency, sphere_dim
    return jnp.where(
        l == 0, 
        jnp.ones_like(l, dtype=jnp.int64), 
        comb(l + d - 2, l - 1) + comb(l + d - 1, l),
    )


# FIXME temporary hack 

from scipy.special import comb as comb_static
from numpy.typing import ArrayLike


def num_phases_in_frequency_static(sphere_dim: int, frequency: ArrayLike) -> ArrayLike:
    l, d = frequency, sphere_dim
    return np.where(
        l == 0, 
        np.ones_like(l, dtype=np.int64), 
        comb_static(l + d - 2, l - 1) + comb_static(l + d - 1, l),
    )


# TODO translate into jax 
def num_phases_to_num_levels(num_phases: int, *, sphere_dim: int) -> int:
    l = 0
    while num_phases > 0:
        num_phases -= num_phases_in_frequency(frequency=l, sphere_dim=sphere_dim)
        l += 1
    return l - 1 if num_phases == 0 else l - 2


def sphere_uniform(sphere_dim: int, n: int, *, key: Key) :
    x = jax.random.normal(key, (n, sphere_dim + 1))
    return x / jnp.linalg.norm(x, axis=-1, keepdims=True)


def sphere_meshgrid(n_colat=100, n_lon=100, epsilon=1e-32):
    colat = np.linspace(0 + epsilon, np.pi - epsilon, n_colat)
    lon = np.linspace(0, 2 * np.pi, n_lon)
    lon, colat = np.meshgrid(lon, colat)
    return colat, lon


def sphere_uniform_grid(n: int) -> Array:
    """
    Fibonacci lattice on the sphere in Cartesian coordinates. 
    """
    phi = (1 + jnp.sqrt(5)) / 2  # Golden ratio
    
    indices = jnp.arange(n)
    theta = 2 * jnp.pi * indices / phi
    phi = jnp.arccos(1 - 2 * (indices + 0.5) / n)
    return sph_to_car(jnp.column_stack((phi, theta)))


hodge_star_matrix = jnp.array([
    [0.0, 1.0],
    [-1.0, 0.0],
])


@Partial(jax.jit, static_argnames=('min_value', ))
def clip_colatitude(x, min_value: float) :
    return jax.lax.cond(
        x[0] < min_value,
        lambda: x.at[0].set(min_value),
        lambda: x,
    )


@jax.jit
def tangent_basis(x) :
    tb = jax.jacfwd(sph_to_car)(x)
    tb /= jnp.linalg.norm(tb, axis=0, keepdims=True)
    return tb 


@jax.jit
def tangent_basis_normalization_matrix(x) :
    return jnp.array([
        [1.0, 0.0], 
        [0.0, 1.0 / jnp.sin(x[0])],
    ])


@Partial(jax.jit, static_argnames=("colatitude_min_value", ))
def expmap_sph(x, v, colatitude_min_value: float = 1e-12) :
    """


Complete the function implementation replacing "raise NotImplementedError".