def assert_compatible_shapes(value_shape, index_shape):
  """Check shapes of the indices and the tensor to be indexed.

  If all input shapes are known statically, obtain shapes of arguments and
  perform compatibility checks. Otherwise, print a warning. The only check
  we cannot perform statically (and do not attempt elsewhere) is making
  sure that each action index in actions is in [0, num_actions).

  Args:
    value_shape: static shape of the values.
    index_shape: static shape of the indices.
  """
  # note: rank-0 "[]" TensorShape is still True.
  if value_shape and index_shape:
    try:
      msg = ("Shapes of \"values\" and \"indices\" do not correspond to "
             "minibatch (2-D) or sequence-minibatch (3-D) indexing")
      assert (value_shape.ndims, index_shape.ndims) in [(2, 1), (3, 2)], msg
      msg = ("\"values\" and \"indices\" have incompatible shapes of {} "
             "and {}, respectively").format(value_shape, index_shape)
      assert value_shape[:-1].is_compatible_with(index_shape), msg
    except AssertionError as e:
      raise ValueError(e)  # Convert AssertionError to ValueError.

  else:  # No shape information is known ahead of time.
    tf.logging.warning(
        "indexing function cannot get shapes for tensors \"values\" and "
        "\"indices\" at construction time, and so can't check that their "
        "shapes are valid or compatible. Incorrect indexing may occur at "
        "runtime without error!")
Check shapes of the indices and the tensor to be indexed.

  If all input shapes are known statically, obtain shapes of arguments and
  perform compatibility checks. Otherwise, print a warning. The only check
  we cannot perform statically (and do not attempt elsewhere) is making
  sure that each action index in actions is in [0, num_actions).

  Args:
    value_shape: static shape of the values.
    index_shape: static shape of the indices.
  
trfl/indexing_ops.py