code: def assert_assigned_type_and_shape_match(existing_tensor, new_tensor):
      assert isinstance(
          new_tensor, jax.Array
      ), f"Assignments to variable must be an instance of JaxArray, but received f{type(new_tensor)}."
  
      new_tensor_device, new_tensor_shape = split_shape_and_device(
          new_tensor
      )
      self_device, self_shape = split_shape_and_device(existing_tensor)
  
      device_mismatch_error = (
          f"Can not replicate a variable that is currently on "
          f"{self_device} devices to {new_tensor_device} devices."
      )
      assert (
          (new_tensor_device is None)
          or (self_device is None)
          or (self_device == new_tensor_device)
      ), device_mismatch_error
  
      shorter_length = min(
          len(new_tensor.shape), len(existing_tensor.shape)
      )
      is_special_ok = isinstance(
          new_tensor, TRACER_TYPES
      ) or isinstance(existing_tensor, TRACER_TYPES)
      is_special_ok = (
          is_special_ok
          and existing_tensor.shape[-shorter_length:]
          == new_tensor.shape[-shorter_length:]
      )
  
      shape_mismatch_error = (
          f"Assign can not change shape of variable. The current variable shape is {self_shape},"
          f" but the requested new shape is {new_tensor_shape}."
      )
      assert (
          is_special_ok
          or new_tensor_shape == self_shape
          or new_tensor.shape == existing_tensor.shape
      ), shape_mismatch_error
  
label_desc: Use image node_id 2 for calling a global function with example usage: assert_assigned_type_and_shape_match(existing_tensor, new_tensor) without return types