tfp.mcmc.*StepSizeAdaptation does not work out of the box with NUTS as it does not have the same previous_kernel_result structure. You will need to manually specify the function to get/set step size, and log_accept_prob:
Oh I see, thanks for the context. It is interesting indeed and I have not thought of that... assuming we can get the gradient of the step_size re log_accept_ratio it should be doable. Not sure if @SiegeLordEx had look into it yet but.