Objective deriv_mode using thing.dim_x results in different Jacobian size

[YigitElma]

The deriv_mode of each objective is decided by comparing the objectives dim_f to the half of the dim_x of the thing it is optimizing. After merging #1052 to the Poincare branch, I realized that the extra Poincare variables switch some objectives from fwd to rev mode auto-diff which is probably ok, but I had to remove jac_chunk_size argument from test_ATF_results since it was triggering one of the errors related to jac_chunk_size and deriv_mode consistencies.

Any thoughts on this?

[YigitElma]

And probably another related issue is that if we have a single objective, the deriv mode being blocked or batched shouldn't cause a difference, right? If so, we may need an additional logic for set_derivatives to use batched since it is more suited for jac_chunk_size's checks.

[dpanici]

hm, what is the error with the ATF one?

[YigitElma]

This one,

errorif(
    self._jac_chunk_size not in ["auto", None]
    and self._deriv_mode == "blocked",
    ValueError,
    "'jac_chunk_size' was passed into ObjectiveFunction, but the "
    "ObjectiveFunction is using 'blocked' deriv_mode, so sub-objective "
    "'jac_chunk_size' are used to compute each sub-objective's Jacobian, "
    "`ignoring the ObjectiveFunction's 'jac_chunk_size'.",
)

The set resolution and added Poincare parameters make the ForceBalance objective have rev mode, and it causes the overall ObjectiveFunction to be deriv_mode=blocked but it is a single objective. Then the given jac_chunk_size to ObjectiveFunction throws the error expecting that sub-objective should get that argument.

I would try to change where to give the jac-chunk_size but Forcebalance is not specified explicitly in that test, but comes from get_equilibrium_objective.

[dpanici]

How does adding poincare parameters cause the mode to change to rev? the dim_x of the objective is only the Rlmn Zlmn lambda LMN, and the dimf is set by the grid size?

Maybe these shuld be warnings though instead of error, I dont like that it can pop up like this

[YigitElma]

No, dim_x is the sum of all optimizable parameters, so R_lmn+Z_lmn+L_lmn+Rb_lmn+Zb_lmn+Rp_lmn + ... axis stuff etc. And the reduced optimization variables are calculated during build of LinearConstraintProjection which doesn't change the dim_x or doesn't call set_derivatives.

[f0uriest]

I think this is because it uses jac_chunk_size=N and deriv_mode="auto" whereas it should be specifying deriv_mode="batched"

[f0uriest]

I see a few different cases:

	deriv_mode="auto"	deriv_mode="batched"	deriv_mode="blocked"
chunk_size=None	no chunking	no chunking	no chunking
chunk_size=N	???	chunking	error
chunk_size="auto"	???	chunking	???

What do we want the behavior to be for the combos marked above?

[YigitElma]

I think this is because it uses jac_chunk_size=N and deriv_mode="auto" whereas it should be specifying deriv_mode="batched"

Yeah, what I mean is if user doesn't specify deriv_mode=batched, it is chosen by Jacobian dimensions (but this is a bit misleading) and if everything is forward mode then deriv_mode=batched. But why not deriv_mode=batched if there is single objective and that objective is reverse mode? We should be able to chunk vjps in the same manner, right?

[f0uriest]

Possible fix:

• chunk_size=N & deriv_mode="auto" -> error, say that specifying chunk size requires batched mode
• chunk_size="auto" & deriv_mode="blocked" -> apply automatic chunk size to sub-objectives that don't have a chunk size already specified?
• chunk_size="auto" & deriv_mode="auto" -> as above, or as batched, depending on what sub-objectives have for deriv_mode

I think this avoids any explicit conflicts, and makes the default of "auto/auto" generally pretty efficient/smart

[dpanici]

I am fine with this comment above