Building a differentiable compute graph dynamically, a la Pytorch

[eguiraud]

Hello,
I was wondering whether in JAX there is a way to dynamically build a differentiable operation, as you can do in Pytorch by operating on tensors.
Something like this (better example below):

x = to.tensor([1., 2., 3.], requires_grad=True)
y = to.tensor([4. ,5. ,6.], requires_grad=True)
z = x * y

if some_dynamic_config_parameter:
  t = z.sum() / 2
else:
  t = z.sum()

t.backward()
print(x.grad, y.grad)

The motivation is that I need to differentiate through computations that are described in config files.
As I parse those config files I can build the computation graph dynamically, but I cannot write a single Python function that represents those computations ahead of time.

I could dynamically produce the code of such a function and eval it, but I was wondering if JAX had a mode of operation more similar to the one in the example above.

I looked around the docs and the GitHub discussions and I could not find an answer to this question -- I'm probably not using the same terms that people in this community would use.

Thanks in advance!
Enrico

[jakevdp]

Thanks for the question! Perhaps I'm misunderstanding, but would something like this work for you?

import jax
import jax.numpy as jnp

def f():
  x = jnp.array([1., 2., 3.])
  y = jnp.array([4., 5., 6.])
  def g(x, y):
    z = x * y

    if some_dynamic_config_parameter:
      t = z.sum() / 2
    else:
      t = z.sum()
    return t
  return jax.grad(g, (0, 1))(x, y)


some_dynamic_config_parameter = True
print(f())
# (Array([2. , 2.5, 3. ], dtype=float32), Array([0.5, 1. , 1.5], dtype=float32))


some_dynamic_config_parameter = False
print(f())
# (Array([4., 5., 6.], dtype=float32), Array([1., 2., 3.], dtype=float32))

In other words, you parse the config file in the definition of the function you'd like to take the gradient of.

[eguiraud]

Thank you for the swift reply! I'm afraid I didn't do a good job at capturing my use case in my simple example 😅 sorry. Let me try again.

Basically I have files that describe math formulae, and code that reads those files and builds a compute graph that corresponds to the formulae. Now I would like to compute the gradient of those formulae. Trying to simplify as much as possible but not too much, in Pytorch I would do it like this:

import torch as to

# Our starting point: an AST that expresses `x*x + y`
# (in reality the AST is built by parsing a math expression coming from an input file)
ast = {"+": [{"*": [{"in": "x"}, {"in": "x"}]}, {"in": "y"}]}


def compute(ast: dict, inputs: dict) -> to.Tensor:
    """Walk the AST, apply the corresponding operations to the inputs"""
    key, value = next(iter(ast.items()))
    if key == "+":
        return compute(value[0], inputs) + compute(value[1], inputs)
    elif key == "*":
        return compute(value[0], inputs) * compute(value[1], inputs)
    elif key == "in":
        return inputs[value]


# Pytorch input tensors
inputs = {
    "x": to.tensor([1.0], requires_grad=True),
    "y": to.tensor([2.0], requires_grad=True),
}

# Apply the operation to Pytorch tensors:
# Pytorch dynamically records what operations are performed on them
# during the forward pass, so it can do reverse mode autodiff
out = compute(ast, inputs)

# Run backprop/reverse mode autodiff
out.backward()

print(f'{inputs["x"].grad=}, {inputs["y"].grad=}')

In JAX the only way I know how to do this is by building a string containing Python code for the function corresponding to the formula, which I then eval:

import jax.numpy as jnp
from jax import grad

# Our starting point: an AST that expresses `x*x + y`
# (in reality the AST is built by parsing a math expression coming from an input file)
ast = {"+": [{"*": [{"in": "x"}, {"in": "x"}]}, {"in": "y"}]}


def synthesize(ast: dict) -> str:
    """Walk the AST and build code for the corresponding Python function"""
    used_inputs = set()

    def write_fn_body(ast):
        key, value = next(iter(ast.items()))
        if key == "+":
            return f"{write_fn_body(value[0])} + {write_fn_body(value[1])}"
        elif key == "*":
            return f"{write_fn_body(value[0])} * {write_fn_body(value[1])}"
        elif key == "in":
            used_inputs.add(value)
            return value

    fn_body = write_fn_body(ast)

    # we need to establish a convention for the ordering of the inputs.
    # alphabetical seems reasonable.
    used_inputs = sorted(used_inputs)
    fn_header = f"lambda {','.join([var for var in used_inputs])}: "

    return eval(fn_header + fn_body), used_inputs


f, used_inputs = synthesize(ast)
df = grad(f, argnums=[0, 1])

inputs = {
    "x": jnp.array(1.0),
    "y": jnp.array(2.0),
}

print(df(*(inputs[k] for k in used_inputs)))

As you see the JAX version is a bit more complex, and since it's code that writes code it's harder to debug and to reason about. So I was wondering if I can do with JAX something similar to what I do in Pytorch in this example.

[jakevdp]

I don't see any reason why you can't just evaluate the expression directly in jax, just as you do in pytorch. For example, this uses the same compute function from your pytorch example:

from functools import partial
import jax.numpy as jnp
import jax

ast = {"+": [{"*": [{"in": "x"}, {"in": "x"}]}, {"in": "y"}]}

def compute(ast: dict, inputs: dict) -> jax.Array:
    """Walk the AST, apply the corresponding operations to the inputs"""
    key, value = next(iter(ast.items()))
    if key == "+":
        return compute(value[0], inputs) + compute(value[1], inputs)
    elif key == "*":
        return compute(value[0], inputs) * compute(value[1], inputs)
    elif key == "in":
        return inputs[value]

inputs = {
    "x": jnp.float32(1.0),
    "y": jnp.float32(2.0),
}
print(jax.grad(partial(compute, ast))(inputs))

{'x': Array(2., dtype=float32), 'y': Array(1., dtype=float32)}

[eguiraud]

Wow, thank you @jakevdp , I just didn't know that was possible!

If I can bother you with a follow-up question just for my understanding...partial(compute, ast) is a weird beast. How does JAX compute its derivative?

[jakevdp]

partial(compute, ast) is just a function that takes some arrays, uses some Python logic to perform some JAX operations on those arrays, and returns the result. JAX can trace it just as it traces any other function that uses Python logic to perform JAX operations on arrays.

[eguiraud]

Makes sense, thank you very much @jakevdp !