Average

Compute the weighted average along the specified axis.

Abstract Signature:

Average(a: Tensor, axis: int, weights: Tensor, returned: bool = False)

PyTorchJAX (Core)NumPyKerasTensorFlowApple MLXFlax NNXPaxML / Praxis

PyTorch

API: torch.mean

Strategy: Direct Mapping

Official Docs ↗

JAX (Core)

API: jax.numpy.average

Strategy: Direct Mapping

Official Docs ↗

NumPy

API: numpy.average

Strategy: Direct Mapping

Official Docs ↗

Keras

API: keras.ops.average

Strategy: Direct Mapping

Official Docs ↗

TensorFlow

API: tf.math.reduce_mean

Strategy: Direct Mapping

Official Docs ↗

Apple MLX

API: mlx.core.mean

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: jax.numpy.average

Strategy: Direct Mapping

Official Docs ↗

PaxML / Praxis

API: jax.numpy.average

Strategy: Direct Mapping

Official Docs ↗