code: def lhs_broadcast(source, target):
      """Ensures that source is compatible with target for broadcasting."""
      same_shape = np.array(source.shape) == np.array(
          target.shape[: source.ndim]
      )
      ones = np.array(source.shape) == np.ones((source.ndim,))
      if np.all(same_shape + ones):
          broadcast_shape = source.shape + (1,) * (
              target.ndim - source.ndim
          )
          return jnp.reshape(source, broadcast_shape)
      raise ValueError(
          f"source shape {source.shape} is not compatible with "
          f"target shape {target.shape}"
      )
  
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 3 for calling a global function with example usage: lhs_broadcast(source, target) and returns: jnp