def to_jax() -> ModuleType:
    """Convert Kornia to JAX.

    Transpiles the Kornia library to JAX using [ivy](https://github.com/ivy-llc/ivy). The transpilation process
    occurs lazily, so the transpilation on a given kornia function/class will only occur when it's called or
    instantiated for the first time. This will make any functions/classes slow when being used for the first time,
    but any subsequent uses should be as fast as expected.

    Return:
        The Kornia library transpiled to JAX

    Example:

    .. highlight:: python
    .. code-block:: python

        import kornia
        jax_kornia = kornia.to_jax()
        import jax
        input = jax.random.normal(jax.random.key(42), shape=(2, 3, 4, 5))
        gray = jax_kornia.color.gray.rgb_to_grayscale(input)

    """
    return ivy.transpile(
        kornia,
        source="torch",
        target="jax",
    )  # type: ignore
Convert Kornia to JAX.

    Transpiles the Kornia library to JAX using [ivy](https://github.com/ivy-llc/ivy). The transpilation process
    occurs lazily, so the transpilation on a given kornia function/class will only occur when it's called or
    instantiated for the first time. This will make any functions/classes slow when being used for the first time,
    but any subsequent uses should be as fast as expected.

    Return:
        The Kornia library transpiled to JAX

    Example:

    .. highlight:: python
    .. code-block:: python

        import kornia
        jax_kornia = kornia.to_jax()
        import jax
        input = jax.random.normal(jax.random.key(42), shape=(2, 3, 4, 5))
        gray = jax_kornia.color.gray.rgb_to_grayscale(input)

    
kornia/transpiler/transpiler.py