# Path: autograd/extend.py

# 
# Path: autograd/numpy/numpy_vjps.py
# def unbroadcast_f(target, f):
#     target_meta = anp.metadata(target)
#     return lambda g: unbroadcast(f(g), target_meta)
# 
# def repeat_to_match_shape(g, shape, dtype, axis, keepdims):
#     """Returns the array g repeated along axis to fit vector space vs.
#        Also returns the number of repetitions of the array."""
#     if shape == ():
#       return g, 1
#     axis = list(axis) if isinstance(axis, tuple) else axis
#     new_shape = onp.array(shape)
#     new_shape[axis] = 1
#     num_reps = onp.prod(onp.array(shape)[axis])
#     # Can't use broadcast_to because of numpy bug: https://github.com/numpy/numpy/issues/9165
#     # return anp.broadcast_to(anp.reshape(g, new_shape), shape), num_reps
#     return anp.reshape(g, new_shape) + onp.zeros(shape, dtype=dtype), num_reps


import scipy.special
import autograd.numpy as np
from autograd.extend import primitive, defvjp, defjvp
from autograd.numpy.numpy_vjps import unbroadcast_f, repeat_to_match_shape
### Error Function ###
inv_root_pi = 0.56418958354775627928
erf = primitive(scipy.special.erf)
erfc = primitive(scipy.special.erfc)

defvjp(erf, lambda ans, x: lambda g:  2.*g*inv_root_pi*np.exp(-x**2))
defvjp(erfc,lambda ans, x: lambda g: -2.*g*inv_root_pi*np.exp(-x**2))


### Inverse error function ###
root_pi = 1.7724538509055159
erfinv = primitive(scipy.special.erfinv)
erfcinv = primitive(scipy.special.erfcinv)

defvjp(erfinv,lambda ans, x: lambda g: g * root_pi / 2 * np.exp(erfinv(x)**2))
defvjp(erfcinv,lambda ans, x: lambda g: -g * root_pi / 2 * np.exp(erfcinv(x)**2))

### Logit and Expit ###
logit = primitive(scipy.special.logit)
expit = primitive(scipy.special.expit)

defvjp(logit,lambda ans, x: lambda g: g / ( x * (1 - x)))
defvjp(expit,lambda ans, x: lambda g: g * ans * (1 - ans))

### logsumexp ###
logsumexp = primitive(scipy.special.logsumexp)

def make_grad_logsumexp(ans, x, axis=None, b=1.0, keepdims=False):
    shape, dtype = np.shape(x), np.result_type(x)
    def vjp(g):

# Path: autograd/extend.py

# 
# Path: autograd/numpy/numpy_vjps.py
# def unbroadcast_f(target, f):
#     target_meta = anp.metadata(target)
#     return lambda g: unbroadcast(f(g), target_meta)
# 
# def repeat_to_match_shape(g, shape, dtype, axis, keepdims):
#     """Returns the array g repeated along axis to fit vector space vs.
#        Also returns the number of repetitions of the array."""
#     if shape == ():
#       return g, 1
#     axis = list(axis) if isinstance(axis, tuple) else axis
#     new_shape = onp.array(shape)
#     new_shape[axis] = 1
#     num_reps = onp.prod(onp.array(shape)[axis])
#     # Can't use broadcast_to because of numpy bug: https://github.com/numpy/numpy/issues/9165
#     # return anp.broadcast_to(anp.reshape(g, new_shape), shape), num_reps
#     return anp.reshape(g, new_shape) + onp.zeros(shape, dtype=dtype), num_reps


# Path: autograd/scipy/special.py
import scipy.special
import autograd.numpy as np
from autograd.extend import primitive, defvjp, defjvp
from autograd.numpy.numpy_vjps import unbroadcast_f, repeat_to_match_shape

### Error Function ###
inv_root_pi = 0.56418958354775627928
erf = primitive(scipy.special.erf)
erfc = primitive(scipy.special.erfc)

defvjp(erf, lambda ans, x: lambda g:  2.*g*inv_root_pi*np.exp(-x**2))
defvjp(erfc,lambda ans, x: lambda g: -2.*g*inv_root_pi*np.exp(-x**2))


### Inverse error function ###
root_pi = 1.7724538509055159
erfinv = primitive(scipy.special.erfinv)
erfcinv = primitive(scipy.special.erfcinv)

defvjp(erfinv,lambda ans, x: lambda g: g * root_pi / 2 * np.exp(erfinv(x)**2))
defvjp(erfcinv,lambda ans, x: lambda g: -g * root_pi / 2 * np.exp(erfcinv(x)**2))

### Logit and Expit ###
logit = primitive(scipy.special.logit)
expit = primitive(scipy.special.expit)

defvjp(logit,lambda ans, x: lambda g: g / ( x * (1 - x)))
defvjp(expit,lambda ans, x: lambda g: g * ans * (1 - ans))

### logsumexp ###
logsumexp = primitive(scipy.special.logsumexp)

def make_grad_logsumexp(ans, x, axis=None, b=1.0, keepdims=False):
    shape, dtype = np.shape(x), np.result_type(x)
    def vjp(g):

        g_repeated,   _ = repeat_to_match_shape(g,   shape, dtype, axis, keepdims)