code: def test_replace_masked(self):
      output = base.replace_masked(
          self.data, self.replacement, self.mask
      )
      expected_output = jnp.array(
          [
              [1, 2, 3, 4, 5, 6],
              [-10, -20, -30, -40, -50, -60],
          ]
      )
      # Test output.
      np.testing.assert_allclose(output, expected_output)
  
file_path: rlax/rlax/_src/base_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import numpy
  from rlax._src import base
label_desc: Use image node_id 2 for calling the ReplaceTest obj's underlying member method code with example usage: obj.test_replace_masked() without return types