custom trace estimation resulting in 2 * expected result

[nahid18]

Hi 👋 ,

I'm not sure if this is the best place for my question but I've been working on this project where we are using the [equinox] (https://github.com/patrick-kidger/equinox) internal API (eqxi) to define a trace primitive for our traceax package, which utilizes the linear operators of lineax.

The issue I'm running into is that our trace estimation implementation is producing results that are 2 * the expected value compared to using jax.numpy.trace.

I suspect that the issue is with the _make_identity function. But I am not sure why.

The issue gets fixed if I replace (in the jvp implementation below) this

t_state = estimator.init(t_key, t_operator)
t_result, _ = eqxi.filter_primitive_bind(_estimate_trace_p, t_key, t_operator, t_state, k, estimator)

with this

t_result = jax.numpy.trace(t_operator.as_matrix())

jvp implementation

@eqxi.filter_primitive_jvp
def _estimate_trace_jvp(primals, tangents):
    key, operator, state, k, estimator = primals
    # t_operator := V
    t_key, t_operator, t_state, t_k, t_estimator = tangents
    jtu.tree_map(_assert_false, (t_key, t_state, t_k, t_estimator))
    del t_key, t_state, t_k, t_estimator

    # primal problem of t = tr(A)
    result, stats = eqxi.filter_primitive_bind(_estimate_trace_p, key, operator, state, k, estimator)
    out = result, stats

    # inner prodct in linear operator space => <A, B> = tr(A @ B)
    # d tr(A) / dA = I
    # t' = <tr'(A), V> = tr(I @ V) = tr(V)
    # tangent problem => tr(V)
    key, t_key = rdm.split(key)
    
    t_operator = jtu.tree_map(eqxi.materialise_zeros, operator, t_operator, is_leaf=_is_none)
    t_operator = lx.TangentLinearOperator(operator, t_operator)
     
    t_state = estimator.init(t_key, t_operator)
    t_result, _ = eqxi.filter_primitive_bind(_estimate_trace_p, t_key, t_operator, t_state, k, estimator)
    
    t_out = (
        t_result,
        jtu.tree_map(lambda _: None, stats),
    )

    return out, t_out

transpose implementation

@eqxi.filter_primitive_transpose(materialise_zeros=True)  # pyright: ignore
def _estimate_trace_transpose(inputs, cts_out):
    # the jacobian, for the trace is just the identity matrix, i.e. J = I
    # so J'v = I v = v

    # primal inputs; operator should have UndefinedPrimal leaves
    key, operator, state, _, estimator = inputs

    # co-tangent of the trace approximation and the stats (None)
    cts_result, _ = cts_out

    # the internals of the operator are UndefinedPrimal leaves so
    # we need to rely on abstract values to pull structure info
    op_t = _make_identity(operator, cts_result)

    key_none = jtu.tree_map(lambda _: None, key)
    state_none = (None, op_t, None)
    k_none = None
    estimator_none = jtu.tree_map(lambda _: None, estimator)

    return key_none, op_t, state_none, k_none, estimator_none

_make_identity implementation

@ft.singledispatch
def _make_identity(op: lx.AbstractLinearOperator, ct_result: float) -> lx.AbstractLinearOperator:
    raise ValueError("Unsupported type!")


@_make_identity.register
def _(op: lx.MatrixLinearOperator, ct_result: float) -> lx.AbstractLinearOperator:
    operator_struct = jtu.tree_map(_remove_undefined_primal, op, is_leaf=_is_undefined)
    in_size = eqx.filter_eval_shape(lambda o: o.in_size(), operator_struct)
    diag = jnp.full(in_size, ct_result)
    return lx.MatrixLinearOperator(jnp.diag(diag), tags=operator_struct.tags)


@_make_identity.register
def _(op: lx.MulLinearOperator, ct_result: float) -> lx.AbstractLinearOperator:
    inner_op = _make_identity(op.operator, ct_result)
    scalar = jnp.array(1.0)
    return lx.MulLinearOperator(inner_op, scalar)


@_make_identity.register
def _(op: lx.TangentLinearOperator, ct_result: float) -> lx.AbstractLinearOperator:
    p_op = op.primal
    t_op = op.tangent
    return lx.TangentLinearOperator(p_op, _make_identity(t_op, ct_result))

minimal test case that demonstrates the issue:

import jax
import jax.numpy as jnp
import jax.random as rdm
import lineax as lx
import traceax as tx

SEED = 0
N = 50
K = 5
key = rdm.PRNGKey(SEED)

def _tr(op, k):
    result = tx.trace(key, op, k)
    return result.value

op = lx.MatrixLinearOperator(jnp.eye(N))
g = jax.grad(_tr)(op, K)
print(f"Grad of mat {op} is {g.as_matrix()}")

The resulting gradient should be an identity matrix, but our code is producing a matrix with all elements equal to 2 on the diagonal.

I am not quite sure what I am missing here.
Your input would be appreciated.
Thank you so much!

[dfm]

Thanks for the question! This doesn't seem like a JAX discussion to me since it builds so much on lineax and equinox. I'd recommend trying to isolate the issue related to one of those packages and open an issue or discussion on one of those repositories. I expect you'll get more useful feedback there, because I don't know enough about the details of those libraries to provide useful suggestions. If you do find that there's a JAX issue once you get to the bottom of it, please report back with the JAX specific example code.