Override __call__ in _Weighting and Aggregator to specify type hints
Force _as_tensor_list to cast tensors into a list of tensors
Fix wrong type provided to _grad and _jac
Rename v to grad_outputs in get_vjp
I think one of the best changes is get_vjp. It was not clear at all what v was supposed to be.
In fact, jac_outputs is a Sequence[Tensor] and v is also a Sequence[Tensor], but of each tensor has 1 less dimension. So v is what we refer to as grad_outputs. I thus renamed v to grad_outputs.
ExceptionContexttype hint for expectations__call__in_WeightingandAggregatorto specify type hints_as_tensor_listto cast tensors into a list of tensors_gradand_jacvtograd_outputsinget_vjpI think one of the best changes is
get_vjp. It was not clear at all whatvwas supposed to be. In fact,jac_outputsis aSequence[Tensor]andvis also aSequence[Tensor], but of each tensor has 1 less dimension. Sovis what we refer to asgrad_outputs. I thus renamedvtograd_outputs.