Hiroki Naganuma

Forward and Backward Mode

0_kgBP3w8p4u0FXicq

Screen Shot 2022-04-13 at 18 27 31

Screen Shot 2022-04-13 at 19 04 42 Screen Shot 2022-04-13 at 19 03 57 Screen Shot 2022-04-13 at 19 03 17 Screen Shot 2022-04-13 at 19 00 45 Screen Shot 2022-04-13 at 19 00 27

Jvp

Screen Shot 2022-04-13 at 20 09 49 Screen Shot 2022-04-13 at 20 10 20

Hvp

# reverse-over-reverse, only works for single arguments
def hvp_revrev(f, primals, tangents):
  x, = primals
  v, = tangents
  return grad(lambda x: jnp.vdot(grad(f)(x), v))(x)
    
# forward-over-reverse
def hvp(f, primals, tangents):
  return jvp(grad(f), primals, tangents)[1]
  
# reverse-over-forward
def hvp_revfwd(f, primals, tangents):
  g = lambda primals: jvp(f, primals, tangents)[1]
  return grad(g)(primals)

Hessian

Screen Shot 2022-04-13 at 19 43 41 Screen Shot 2022-04-13 at 19 32 46

Papers