# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_jaxify.ipynb.

# %% auto 0
__all__ = ['jaxit']

# %% ../nbs/01_jaxify.ipynb 2
from ceviche import jacobian
import numpy as np
import jax
import jax.numpy as jnp
from typing import List

# %% ../nbs/01_jaxify.ipynb 3
def jaxit(
  mode: str='reverse', #the mode used to calculate the jacobian using `ceviche`
  argnums: List[int]=[0] #the argument indices this function should be differentiable against
  ):
  """
    make a function that internally uses autograd compatible to jax gradient calculations

    Attention: only a single output variable is supported
  """
  def inner(function):
    grad_fns = [jacobian(function, mode=mode, argnum=i) for i in argnums]

    @jax.custom_jvp
    def jaxed(*args):
      return function(*args)
    
    @jaxed.defjvp
    def jaxed_jvp(primals, tangents):
      primals_out = jaxed(*primals)
      as_np = [np.asarray(prim) for prim in primals]
      grads = [jnp.array(grad_fns[i](*as_np)) for i in argnums]

      # if len(tangents) > len(grads):
      #   raise RuntimeError("passed `num_args` is lower than the actual number of arguments")
      
      contributions = jnp.array([jnp.dot(grads[i],tangents[i].flatten()) for i in argnums])
  
      return primals_out, jnp.sum(contributions)
    return jaxed
  return inner

