Differentiating with respect to closure introduced variables.

[VIVelev]

So, I've got the following example:

import jax
import jax.numpy as jnp
import numpy as np


def Rx(theta):
    """Rotation around the x-axis by angle theta"""

    ct = jnp.cos(theta)
    st = jnp.sin(theta)

    def f(x, y, z):
        return jnp.array([x, ct * y - st * z, st * y + ct * z])

    return f


Rx(np.pi / 2)(
    0.0, 1.0, 0.0
)  # => Array([ 0.000000e+00, -4.371139e-08,  1.000000e+00], dtype=float32)

# Take the derivative of rotation around the x-axis w.r.t. theta at [0, 1, 0]:
jax.jacfwd(Rx)(np.pi / 2)(
    jnp.array([0.0, 1.0, 0.0])
)  # => TypeError: Value <function Rx.<locals>.f at 0x127a16520> with type <class 'function'> is not a valid JAX type

Can I make this work? I would like to differentiate the rotation w.r.t. theta, or in general, apply the jacobian function transformation to functions that take in jnp.Arrays and return a function that performs a computation with those jax.Arrays

[ASEM000]

The problem is that Rx returns a function not an inexact value ( jax array of float or complex or python float)

I slightly modified the code to make it work

import jax
import jax.numpy as jnp


def Rx(theta, x, y, z):
    """Rotation around the x-axis by angle theta"""

    ct = jnp.cos(theta)
    st = jnp.sin(theta)

    out = jnp.array([x, ct * y - st * z, st * y + ct * z])

    return out


print(Rx(jnp.pi / 2, 0.0, 1.0, 0.0))
# => Array([ 0.000000e+00, -4.371139e-08,  1.000000e+00], dtype=float32)

# Take the derivative of rotation around the x-axis w.r.t. theta (argnums=0) at [0, 1, 0]:
jax.jacfwd(Rx, argnums=0)(jnp.pi / 2, 0.0, 1.0, 0)

Let me know if this helps

[VIVelev]

Thanks for the response! Yeah, I thought of that, but was looking for a different workaround. I kinda wanted to describe the parametric family of transformations which are rotations around the x-axis, theta being the parameter, and diff w.r.t. that parameter. A function returning a function is the most concise way to do it, but I guess I shouldnt be so picky.

[jakevdp]

If you don't want to restructure your closure definitions, you can compute the derivative like this:

result = jax.jacfwd(lambda theta: Rx(theta)(0.0, 1.0, 0.0))(np.pi / 2)
print(result)
# [ 0.000000e+00 -1.000000e+00 -4.371139e-08]

There's no built-in JAX transform to do autodiff of a function that returns a function, but you could use a short wrapper to define this behavior if you wish:

def fun_jacfwd(fun, *args, **kwargs):
  def df(*df_args, **df_kwargs):
    return jax.jacfwd(lambda *a: fun(*a)(*df_args, **df_kwargs), **kwargs)(*args)
  return df

Rx_prime = fun_jacfwd(Rx, np.pi / 2)
print(Rx_prime(0.0, 1.0, 0.0))
# [-0.000000e+00 -1.000000e+00 -4.371139e-08]

[VIVelev]

Exactly what I needed. Thanks!!