Pseudocode for discussion on general structure for using jax.value_and_grad: question on general structure, restrictions and best practices

[TomClarkMassSpec]

[jakevdp]

I don't understand your pseudocode. For example, what operation is func_B(x,y) = Func_B(pytree) supposed to represent?