def _select_and_scatter_translation(ctx, avals_in, avals_out, operand, source,
init_value, *, select_jaxpr, select_consts,
scatter_jaxpr, scatter_consts,
window_dimensions, window_strides,
padding):
select = lax._reduction_computation(ctx, select_jaxpr, select_consts,
init_value)
scatter = lax._reduction_computation(ctx, scatter_jaxpr, scatter_consts,
init_value)
return [
xops.SelectAndScatterWithGeneralPadding(operand, select,
window_dimensions,
window_strides, padding,
source, init_value, scatter)
]
def _reduce_window_translation_rule(ctx, avals_in, avals_out, *args, jaxpr,
consts, window_dimensions, window_strides,
padding, base_dilation, window_dilation):
operands, init_values = util.split_list(args, [len(args) // 2])
xla_computation = lax._reduction_computation(ctx,
jaxpr,
consts,
init_values,
singleton=False)
return xla.xla_destructure(
ctx.builder,
xops.ReduceWindowWithGeneralPadding(operands, init_values,
xla_computation, window_dimensions,
window_strides, base_dilation,
window_dilation, padding))