How to represent the jacobian of a function where the domain field is from a cartesian product

autograd
1

Normally, if I have a function f = R^2 → R^3. It will have a 3 by 2 jacobian. But what if my function takes two arguments in R^2.

// pseudocode
f(a, b): // a and b are both (2,1) so this function has 4 parameters
return torch.cat((a + b, [1]), 1) // output is (3,1)

This is the same as saying this is a function of R^4 to R^3.

What would be the dimension of the jacobian here? 3 by 4? 3 by 2 by 2?

2

Hi Cedric!

Pytorch’s approach (which seems fine to me) is to return a tuple of jacobians,
one for each of the arguments. Consider:

>>> import torch
>>> torch.__version__
'2.5.1'
>>> 
>>> _ = torch.manual_seed (2024)
>>> 
>>> def f (r2a, r2b):    # some r2, r2 --> r3 function
...     m = torch.arange (6.0).reshape (3, 2)
...     return  m @ (r2a * r2b)
... 
>>> a = torch.randn (2)
>>> b = torch.randn (2)
>>> 
>>> jacfunc = torch.func.jacrev (f, argnums = (0, 1))   # tell jacrev you want both arguments of f
>>> 
>>> jacfunc (a, b)                                      # jacfunc returns a tuple of both jacobians
(tensor([[-0.0000,  1.3722],
        [-1.6281,  4.1165],
        [-3.2562,  6.8608]]), tensor([[-0.0000,  1.7260],
        [-0.0809,  5.1780],
        [-0.1617,  8.6300]]))

Best.

K. Frank