jax symbolic autodiff on sin cos ...

[jakubMitura14]

Hello I was reading through the autodiff cookbook but one thing is not clear to me, do simple functions like trig functions are differentiated symbolically or using numerical approximations?
so for example when I have function

def a(x):
    return jnp.sin(x)

the jax during reverse autodiff will get a derivative as just cos?

[jakevdp]

JAX does symbolic automatic differentiation at the primitive level. You can see this by printing the jaxpr for the transformed function:

import jax
import jax.numpy as jnp

def a(x):
    return jnp.sin(x)

x = jnp.float32(1.0)

print(jax.make_jaxpr(a)(x))
# { lambda ; a:f32[]. let b:f32[] = sin a in (b,) }

print(jax.make_jaxpr(jax.grad(a))(x))
# { lambda ; a:f32[]. let
#     _:f32[] = sin a
#     b:f32[] = cos a
#     c:f32[] = mul 1.0 b
#   in (c,) }