Compute the mean of the valid entries of x along one axis, where mask marks the valid entries, while keeping every intermediate shape independent of the mask values.
For each slice of x along axis, the result is
result=∣V∣1j∈V∑xjwhere j runs over the positions along axis, V is the set of positions where mask is True, and |V| is the number of such positions. A slice with no valid entries gives 0.0. Entries where mask is False are padding, and their values must have no effect on the result.
Your function is called under jax.jit with axis marked as static. That means x and mask arrive as traced arrays whose values are unknown while the function compiles, and axis arrives as an ordinary Python integer.
Return a float32 JAX array whose shape is the shape of x with axis removed.
Input: x = [[1.0,2.0,3.0,4.0],[10.0,20.0,30.0,40.0]], mask = [[true,true,false,false],[true,false,true,false]], axis = 1
Output: [1.5,20.0]
Explanation: Row 0 keeps 1 and 2, whose mean is 1.5. Row 1 keeps 10 and 30, whose mean is 20.
Input: x = [[1.0,2.0],[3.0,4.0],[5.0,6.0]], mask = [[true,false],[true,true],[false,true]], axis = 0
Output: [2.0,5.0]
Input: x = [[5.0,7.0,9.0],[2.0,4.0,6.0]], mask = [[false,false,false],[true,true,false]], axis = -1
Output: [0.0,3.0]
jnp.where(mask, x, 0.0) keeps the full shape and zeroes the padding.
jnp.sum(mask, axis=axis) counts the valid entries of each slice.
jnp.maximum(count, 1) avoids dividing by zero.
Accepts: array
Accepts: array
Accepts: integer
Compute the mean of the valid entries of x along one axis, where mask marks the valid entries, while keeping every intermediate shape independent of the mask values.
For each slice of x along axis, the result is
result=∣V∣1j∈V∑xjwhere j runs over the positions along axis, V is the set of positions where mask is True, and |V| is the number of such positions. A slice with no valid entries gives 0.0. Entries where mask is False are padding, and their values must have no effect on the result.
Your function is called under jax.jit with axis marked as static. That means x and mask arrive as traced arrays whose values are unknown while the function compiles, and axis arrives as an ordinary Python integer.
Return a float32 JAX array whose shape is the shape of x with axis removed.
Input: x = [[1.0,2.0,3.0,4.0],[10.0,20.0,30.0,40.0]], mask = [[true,true,false,false],[true,false,true,false]], axis = 1
Output: [1.5,20.0]
Explanation: Row 0 keeps 1 and 2, whose mean is 1.5. Row 1 keeps 10 and 30, whose mean is 20.
Input: x = [[1.0,2.0],[3.0,4.0],[5.0,6.0]], mask = [[true,false],[true,true],[false,true]], axis = 0
Output: [2.0,5.0]
Input: x = [[5.0,7.0,9.0],[2.0,4.0,6.0]], mask = [[false,false,false],[true,true,false]], axis = -1
Output: [0.0,3.0]
jnp.where(mask, x, 0.0) keeps the full shape and zeroes the padding.
jnp.sum(mask, axis=axis) counts the valid entries of each slice.
jnp.maximum(count, 1) avoids dividing by zero.
Accepts: array
Accepts: array
Accepts: integer
