wehub-resource-sync
22 lines
613 B
Python
Executable File
22 lines
613 B
Python
Executable File
from prml.nn.math.sum import sum
|
|
|
|
|
|
def mean(x, axis=None, keepdims=False):
|
|
"""
|
|
returns arithmetic mean of the elements along given axis
|
|
"""
|
|
if axis is None:
|
|
return sum(x, axis=None, keepdims=keepdims) / x.size
|
|
elif isinstance(axis, int):
|
|
N = x.shape[axis]
|
|
return sum(x, axis=axis, keepdims=keepdims) / N
|
|
elif isinstance(axis, tuple):
|
|
N = 1
|
|
for ax in axis:
|
|
N *= x.shape[ax]
|
|
return sum(x, axis=axis, keepdims=keepdims) / N
|
|
else:
|
|
raise TypeError(
|
|
"Unsupported type for axis: {}".format(type(axis))
|
|
)
|