code: def replace_masked(
      data: chex.Array,
      replacement: Optional[chex.Array],
      mask: chex.Array,
  ):
      """Replace slices of an array as indicated by a mask.
  
      Args:
        data: an array whose some elements we want to replace.
        replacement: an array with the same shape as `data`, containing the data to
          insert according to `mask`. If `None` is passed, then the masked elements
          in `data` will be replaced with zeros.
        mask: a mask of 0/1s, whose shape is a prefix of `data` and `replacements`.
          When the mask is 1, values of data are replaced.
  
      Returns:
        the updated tensor.
      """
      if replacement is None:
          replacement = jnp.zeros_like(data)
      return jnp.where(lhs_broadcast(mask, data), replacement, data)
  
file_path: rlax/rlax/_src/base.py
imports: from typing import Optional, Sequence, Union
  import chex
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: replace_masked(data, replacement, mask) and returns: jnp