def _get_tensor_dim(tensor: tf.Tensor, dim: int,
                    context: _Context) -> Union[int, tf.Tensor]:
  """Gets  specified tensor dimension, backing off to a bound if ragged."""

  # TODO(jhclark): Add tutorial section on how this backoff works.
  if isinstance(tensor, tf.RaggedTensor):
    bounding_shape = tensor.bounding_shape()
    # We back off to the bounding box for finding named dimensions in
    # `RaggedTensor`s.
    return bounding_shape[dim]

  if tensor.shape.num_elements() is None:
    # Shape is fully unknown.
    dyn_shape = tf.shape(tensor)
    return dyn_shape[dim]

  shape = tensor.shape.as_list()
  if dim >= len(shape):
    raise ValueError(
        "Building contract for {}: Dimension index {} is out of range for "
        "shape ({}) in tensor {}".format(context.func_name, dim, shape,
                                         tensor.name))
  if shape[dim] is not None:
    return shape[dim]
  else:
    dyn_shape = tf.shape(tensor)
    return dyn_shape[dim]
Gets  specified tensor dimension, backing off to a bound if ragged.
language/canine/tensor_contracts.py