Feature request: Add support for multi-argument functions.

[FlyingFaller]

Thanks for building this! Looking very promising if performant and reliable. Would love to see support multiple arguments in the input function. In JAX this has been accomplished with the argnums argument specifying the index or sequence of indices to differentiate with respect to.

[adrhill]

Hi @FlyingFaller, thanks for the issue. I'm open to this idea.

In JAX this has been accomplished with the argnums argument specifying the index or sequence of indices to differentiate with respect to.

Could you provide me an example of this? The lack of a clear input enumeration on multi-arg functions is exactly what has stopped me from implementing this until now.

And out of curiosity: do you have a specific use-case in mind?

[FlyingFaller]

Could you provide me an example of this?

Sure. JAX does AD with the jax.grad(fun, argnums=0, has_aux=False, holomorphic=False, allow_int=False, reduce_axes=()) function where argnums is an int or tuple of ints representing the indices of arguments in fun to differentiate with respect to. Calling jax.grad(fun) on a multi-argument function will return a callable that expects the same number and shape of arguments as fun. jax.grad(fun, argnums=tuple) also returns a single jacobian function, but this callable, in turn, returns a corresponding tuple of jacobians corresponding to argnums. As a more concrete example:

import jax

# simple multi-argument function
def f(x, y):
    return 3.0 * x**2 + 2.0 * y**3

# Default case argnums=0
# Jacobian wrt x evaluated at (x_val, y_val)
grad_x_fn = jax.grad(f, argnums=0) 
grad_x = grad_x_fn(x_val, y_val) 

# Jacobian wrt y evaluated at (x_val, y_val)
grad_y_fn = jax.grad(f, argnums=1)
grad_y = grad_y_fn(x_val, y_val)

# Jacobian wrt both x and y evaluated at (x_val, y_val)
grad_both_fn = jax.grad(f, argnums=(0, 1))
grad_x, grad_y = grad_both_fn(x_val, y_val) # Tuple of Jacobians

For asdex you will probably need to adjust the input_shape argument to accept a tuple/list of shapes corresponding to each argument in the input function. Yesterday I forked the repository and (with a lot of help from LLMs) implemented this (very poorly) for Jacobians, but understanding ASD enough to make it work well is outside my wheelhouse.

And out of curiosity: do you have a specific use-case in mind?

I am doing a lot of work with optimal control theory and trajectory optimization and would really like a good AD engine for linearizing user-made constraint and dynamics functions. My current code is structured in such a way where multi-argument functions make sense. Other libraries like https://github.com/PTNobel/AutoDiff have also done similar stuff.

[adrhill]

Thanks for your patience regarding this feature @FlyingFaller. I originally planned to include this in the breaking v0.2.0 release via #105, but I'm not too happy with the current state of the PR, as Claude Opus 4.7 is struggling to write decent code.

I'm still unsure regarding some aspects of the API design: Likely, I'm going to have to drop the current input_shape argument in favor or users actually passing sample inputs. Python's tacked-on type system is unfortunately fairly limited. An other minor issue is how we represent colored patterns. The unicode printouts, plots (#97) and save/load functionality all have to be rewritten. But we'll get there!

[FlyingFaller]

I think its awesome that you're working on this! One solution would be to pass sample inputs and forgo any kind of input_shape(s). I assume these would be dummy inputs that could be zeros or Nones and they would just be for shape and sparsity detection and would not be evaluated?

I agree your current approach with both input_shape and input_shapes and pytrees feels a bit awkward. I think what I did when I was tinkering was force the way shapes are defined as input_shape: tuple[int, ...] | tuple[tuple[int, ...], ...] so a "shape" always meant a tuple of ints. i.e. (1, 2, 3) is treated as a single 3d input arg and is normalized to ((1,2,3),) and ((4, 5), (6),) is a multi-input arg function. Another way I considered handling this was creating a Shape object or using one of JAX's Shape objects to eliminate any ambiguity. If this commentary is not relevant to the friction you're having feel free to disregard it.

[jpbrodrick89]

Just took a look at this, and wondered if following lineax and using pytree's of ShapeDTypeStruct's would work tracking shape if I understand correctly?

[adrhill]

Yes, I think asking for sample inputs in asdex.jacobian(f, *args) and internally using ShapeDTypeStruct is the way to go.