Understanding linear_transpose

[gitcommitypos]

What exactly does the linear_transpose function do? In the example given in the docs

f = lambda x, y: 0.5 * x - 0.5 * y
scalar = types.SimpleNamespace(shape=(), dtype=np.dtype(np.float32))
f_transpose = jax.linear_transpose(f, scalar, scalar)
print(f_transpose(10.0))

What exactly does the result mean? Is it the inverse of the function?

[YouJiacheng]

Transpose a function that is promised to be linear.
Namely, given a linear function f, it can be written in matrix form: f(x) == A @ x, then jax.linear_transpose(f, *place_holder_args)(y) == A.T @ y

[gitcommitypos]

Thank you very much @YouJiacheng . I also have a related question about linearize. As far as I understand, linearizing a function is

f_lin(x,y) = f(a,b) + (x-a) df/da + (y-b) df/db

In the docs , they say that linearize and jvp are equivalent as

y, out_tangent = jax.jvp(fun, x, r)

y, f_ = jax.linearize(fun, *x)

out_tangent = f_(r)

How can the two provide the same numerical result when $f_lin$ computes extra terms like f(a,b) - a * df/da - b* df/db.

Thanks for your help.

[YouJiacheng]

It is not the case.
For y, f_ = jax.linearize(f, a, b), f_ is not the f_lin that you defined.
Actually f_lin(x, y) == y + f_(x - a, y - b) == f(a,b) + (x-a) df/da + (y-b) df/db
Namely, f_ is a linear function which must not have any constant.
f_lin is usually not a linear function since it contain a constant term f(a,b) - a * df/da - b* df/db.

[gitcommitypos]

Thank you for clarifying.