Optimising a gate with two parameters

[lm1909]

I want to optimise the parameters in a tensor network with autodiff. The tensor network is contructed from gates like circuit.apply_gate('RZ', shared_variable * local_constant, ..., parametrize=parametrize). What I would like to do is to optimise shared_variable for multiple gates simultaneously (in the spirit of shared_tags), while local_constant however, should not be optimised and is a constant specific to the gate. I'm using TNOptimiser for the optimisation. Naively putting the tag of the the tensor node in shared_tags does not work of course, as then all gates will have the same parameter, independent of local_constant.

What would be the best way to only optimise shared_variable and keep it the same over all shared_tags? (I guess another way to phrase the question would be, of how to construct a gate with more than one parameter, where the parameters can be optimised independently)

Thanks a lot in advance!

[lm1909]

Actually, I figured out that this is quite simple to workaround - by simply defining a PArray with a custom function. Conceptually something like:

def helper_function(params):
   return qtn.circuit.rz_gate_param_gen([params[0] * local_constant]])
circuit.apply_gate_raw(qtn.array_ops.PArray(helper_function, [shared_variable]), ..., parametrize=True)