jax.lax.reduce_window#

jax.lax.reduce_window(operand,init_value,computation,window_dimensions,window_strides=None,padding='VALID',base_dilation=None,window_dilation=None)[source]#

Reduction over padded windows.

Wraps XLA’sReduceWindowWithGeneralPadding operator.

Parameters:

• operand (Any) – input array or tree of arrays.

• init_value (Any) – value or tree of values. Tree structure must match thatofoperand.

• computation (Callable) – callable function over which to reduce. Input and output must bea tree of the same structure asoperand.

• window_dimensions (core.Shape) – sequence of integers specifying the window size.

• window_strides (Sequence[int]|None) – optional sequence of integers specifying the strides, ofthe same length aswindow_dimensions. Default (None) indicatesa unit stride in each window dimension.

• padding (str |Sequence[tuple[int,int]]) – string or sequence of integer tuples specifying the type of paddingto use (default: “VALID”). If a string, must be one of “VALID”, “SAME”, or“SAME_LOWER”. See thejax.lax.padtype_to_pads() utility.

• base_dilation (Sequence[int]|None) – optional sequence of integers for base dilation values, ofthe same length aswindow_dimensions. Default (None) indicates unitdilation in each window dimension.

• window_dilation (Sequence[int]|None) – optional sequence of integers for window dilation values, ofthe same length aswindow_dimensions. Default (None) indicates unitdilation in each window dimension.

Returns:

A tree of arrays with the same structure asoperand.

Return type:

Any

Example

Here is a simple example of a windowed product over pairs in a 1-dimensional array:

>>>importjax>>>x=jax.numpy.arange(10,dtype='float32')>>>xArray([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.], dtype=float32)

>>>initial=jax.numpy.float32(1)>>>jax.lax.reduce_window(x,initial,jax.lax.mul,window_dimensions=(2,))Array([ 0.,  2.,  6., 12., 20., 30., 42., 56., 72.], dtype=float32)