¿Cuál es la alternativa de elementwise_grad de autograd en JAX?
Quiero resolver una ecuación diferencial de segundo orden con una red neuronal. Para la diferenciación automática estoy utilizando la biblioteca JAX. Para calcular la primera y segunda derivada de mi variable objetivo ‘u’, es decir, calcular du/dx y d2u/dx2, se ha utilizado elementwise_grad en un ejemplo. ¿Cuál es su alternativa . . . Read more