code: def weighted_max_pool(
      inputs,
      weights,
      window_shape,
      strides=None,
      padding="VALID",
      return_pooled_weights=False,
  ):
      """Pools the input by taking max over a window, w.r.t their inputs' weights.
  
      Args:
        inputs: Input data with dimensions (batch, <window dims>, features).
        weights: Input weights with dimensions (batch, <window dims>).
        window_shape: tuple; A shape tuple defining the window to reduce over.
        strides: tuple; A sequence of `n` integers, representing the inter-window
            strides (default: `(1, ..., 1)`).
        padding: str/list(tuple); Either the string `'SAME'`, the string `'VALID'`,
          or a sequence of `n` `(low, high)` integer pairs that give the padding to
          apply before and after each spatial dimension (default: `'VALID'`).
        return_pooled_weights: bool; Also return the pooled weight
  
      Returns:
        The maximum of each window slice. If return_pooled_weights is True, it also
        returns the maximum of pooled weights.
      """
      assert inputs.shape[:-1] == weights.shape
      weights = jnp.expand_dims(weights, -1)
      inputs = inputs * weights
      outputs = nn.max_pool(
          inputs, window_shape, strides=strides, padding=padding
      )
      if return_pooled_weights:
          max_weights = nn.max_pool(
              weights, window_shape, strides=strides, padding=padding
          )
          return outputs, max_weights.squeeze(axis=-1)
      return outputs
  
file_path: scenic/scenic/model_lib/layers/nn_ops.py
imports: from typing import Callable, Sequence
  import flax.linen
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 7 for calling a global function with example usage: weighted_max_pool(inputs, weights, window_shape, strides, padding, return_pooled_weights) and returns: outputs, outputs,  max_weights