guidelines on how to define the compute graph -- a.k.a., the forward pass

[wsunid]

Hi,

I'm playing your sparse optimizer and meet an error.

Basically, in my compute graph, I calculate two losses from two sets of 2D tracks -- residual A^{m, 2} and residual B^{n, 2} and then concatenate them into a final loss tensor of size (m+n)x2.
But it reports the errors below:

bae/optim/optimizer.py", line 12, in jacobian

assert output.optrace[id(output)][0] == 'map', "The last operation in compute graph being indexing transform is not meaningful"

AttributeError: 'TrackingTensor' object has no attribute 'optrace'. Did you mean: 'trace'? 

It seems that the optimizer has the very strict requirements on how we define the compute graph. Could you please give me some insight regarding why this happen and the guidance on how we should define the proper compute graph?

Thanks,
Weiwei.

[zitongzhan]

I see the computation graph has two residual vectors, which is not common in BA problems, but I think we could still manage to support this. That is, the compute graph traversal (i.e., the jacobian and backward function) needs to be expanded to support more complex compute graph. May I know what's the shape of optimizable parameter(s) and the shape of the Jacobian?

[wsunid]

I think my optimization parameters are a bit complex.

For example, I have cameras A (cx7), points B (nx3) and points C (mx3). The reason I have two point sets is because I need to transform points C with other optimizable transformation parameters (mx7).

Can you provide some quick guideline here?

Thanks!

[zitongzhan]

It is completely possible to build and solve a compute graph like this. We need to add support for torch.cat(..., dim=0) type of operation and treat it as an indexing type in our graph search.

[randomze]

I'm also in need of being able to use torch.concatenate and torch.stack, and maybe some other operations. Where can I find which operations are supported in BAE's operation graph?

[zitongzhan]

The progress of torch.cat support is currently in construction in the cat_thr_slicing_jac branch

[zitongzhan]

@wsunid I finished an initial draft supporting the torch.cat here #11 . Would you be interested in testing it out?

[wsunid]

Thanks Zitong!

I quickly tested and found the issues when we try to retrieve the diagonal value of A matrices for damping. Specifically, at the line: https://github.com/zitongzhan/bae/blob/release/bae/sparse/py_ops.py#L140 , the A matrix has the empty diagonal values, suggesting that some parameters were not linked to observations. This is not true in my case.

So I would recommend having the code tested on typical BA toy data and you may find the issues. My guess is that there might be some issues when you try to trace jacobian in _slice_upstream_tuple_columns, causing the wrong jacobian filling. Though I'm not 100% sure.

Thanks again for the PR.

All the best,
Weiwei.

[zitongzhan]

Thanks Zitong!

I quickly tested and found the issues when we try to retrieve the diagonal value of A matrices for damping. Specifically, at the line: https://github.com/zitongzhan/bae/blob/release/bae/sparse/py_ops.py#L140 , the A matrix has the empty diagonal values, suggesting that some parameters were not linked to observations. This is not true in my case.

So I would recommend having the code tested on typical BA toy data and you may find the issues. My guess is that there might be some issues when you try to trace jacobian in _slice_upstream_tuple_columns, causing the wrong jacobian filling. Though I'm not 100% sure.

Thanks again for the PR.

All the best, Weiwei.

Thanks @wsunid for the quick test!

I added an example that matches your use case (two point sets where one subset is transformed by additional optimizable SE(3) parameters, then combined with torch.cat(..., dim=0)), and added checks that (1) the combined COO Jacobian has no empty columns and (2) diag(JᵀJ) is fully occupied (no “missing” diagonal entries for damping).

I also ran the existing “real” BAL Jacobian tests (using the standard BAL datasets) and they do not expose the issue on their own. The new "split-points + transform + cat" test here is what reproduces the structure you described and guards against regressions.

Also, if you’re able to share your exact loss function (or a minimal runnable snippet that defines it), I’m happy to plug it in and try to reproduce the original empty-diagonal behavior directly. That would help confirm whether there are any remaining edge cases beyond the BAL-style cat example.

[zitongzhan]

Thanks for asking this @randomze and sorry for the late reply

The best place to look is the “What each tracked op means” I just added: https://github.com/sair-lab/bae/blob/release/.agent/skills/bae-compute-graph/SKILL.md#what-each-tracked-op-means

At the operation-graph level, BAE currently special-cases three tracked op types: index, map, and torch.cat(..., dim=0). So torch.concatenate is only supported if it is effectively torch.cat(..., dim=0), while torch.stack is not currently a tracked structural op. If an operation is used inside @map_transform, the rule is broader: it can work as long as torch.func.jacrev/vmap can differentiate it. The implementation-side source of truth is bae/autograd/function.py and bae/autograd/graph.py.

If you need more operation support just let me know to add them.

[zitongzhan]

Hi @wsunid

I checked this case on our side and added a BAL-based regression test around cat inside the gauge-fixed setup.

In particular, I tested a “fix the first camera, optimize the rest” formulation where the forward pass explicitly concatenates the fixed camera pose with the remaining camera poses before projection:

camera_se3 = torch.cat([camera_fixed, pose_rest], dim=0)
points_proj = project_with_se3_and_intrinsics(
    points_3d[point_indices],
    camera_se3[camera_indices],
    intrinsics[camera_indices],
)
residual = points_proj - points_2d

I then validated two things:

1. The Jacobian structure is still correct in the gauge-fixed case after the cat.

   • after concatenating these Jacobians, there are no empty columns and diag(J^T J) is fully occupied

2. The gauge-fixed BA reproduces the same optimization result as the non-gauge-fixed formulation.

   • I compared the final per-pixel error against the standard BAL test path and got matching results within tolerance

So at least for this BAL gauge-fixing pattern, using torch.cat(...) in the forward pass behaves correctly.

The concrete test is in tests/autograd/test_bal_jacobian.py, under test_bal_jacobian_cat_fixed_first_camera_gauge_free(...), if you want to see the exact setup.

Let me know if you need anything

[wsunid]

Hi Zitong,

Thanks for the fixes! I'll test the fix and will let you know soon.

All the best,
Weiwei.