def weighted_average(
F, x: Tensor, weights: Optional[Tensor] = None, axis: Optional[int] = None
) -> Tensor:
"""
Computes the weighted average of a given tensor across a given axis, masking values associated with weight zero,
meaning instead of `nan * 0 = nan` you will get `0 * 0 = 0`.
Parameters
----------
F
The function space to use.
x
Input tensor, of which the average must be computed.
weights
Weights tensor, of the same shape as `x`.
axis
The axis along which to average `x`
Returns
-------
Tensor:
The tensor with values averaged along the specified `axis`.
"""
if weights is not None:
weighted_tensor = F.where(
condition=weights, x=x * weights, y=F.zeros_like(x)
)
sum_weights = F.maximum(1.0, weights.sum(axis=axis))
return weighted_tensor.sum(axis=axis) / sum_weights
else:
return x.mean(axis=axis)
def weighted_average(F,
x: Tensor,
weights: Optional[Tensor] = None,
axis=None) -> Tensor:
"""
Computes the weighted average of a given tensor across a given axis.
Parameters
----------
F
The function space to use.
x
Input tensor, of which the average must be computed.
weights
Weights tensor, of the same shape as `x`.
axis
The axis along which to average `x`
Returns
-------
Tensor:
The tensor with values averaged along the specified `axis`.
"""
if weights is not None:
weighted_tensor = x * weights
sum_weights = F.maximum(1.0, weights.sum(axis=axis))
return weighted_tensor.sum(axis=axis) / sum_weights
else:
return x.mean(axis=axis)
def weighted_average(
F, tensor: Tensor, weights: Optional[Tensor] = None, axis=None
):
if weights is not None:
weighted_tensor = tensor * weights
sum_weights = F.maximum(1.0, weights.sum(axis=axis))
return weighted_tensor.sum(axis=axis) / sum_weights
else:
return tensor.mean(axis=axis)
