class_docstr: Tests for utilities in nn_ops.py.
class_docstr_tok: ["Tests","for","utilities","in","nn_ops.py","."]
docstr: Tests weighted_max_pool.
docstr_tok: ["Tests","weighted_max_pool","."]
code: def test_weighted_max_pool(self):
      """Tests weighted_max_pool."""
      inputs_shape = (16, 32, 32, 20)
      window_shape = (4, 4)
      strides = (4, 4)
      inputs = jnp.array(np.random.normal(size=inputs_shape))
      weights = jnp.ones(inputs_shape[:-1])
  
      outputs, pooled_weights = nn_ops.weighted_max_pool(
          inputs,
          weights,
          window_shape=window_shape,
          strides=strides,
          padding="VALID",
          return_pooled_weights=True,
      )
  
      expected_outputs = nn.max_pool(
          inputs,
          window_shape=window_shape,
          strides=strides,
          padding="VALID",
      )
      expected_pooled_weights = jnp.ones((16, 8, 8))
      self.assertTrue(jnp.array_equal(outputs, expected_outputs))
      self.assertTrue(
          jnp.array_equal(pooled_weights, expected_pooled_weights)
      )
  
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import flax.linen
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import nn_ops
label_desc: Use image node_id 2 for calling the NNOpsTest obj's underlying member method code with example usage: obj.test_weighted_max_pool() without return types