def _bounding_shape(tensor: Union[tf.Tensor, tf.RaggedTensor]) -> tf.Tensor:
  """Returns the bounding shape of a Tensor, which is possibly Ragged."""
  if isinstance(tensor, tf.RaggedTensor):
    # TODO(jhclark): How do we handle dynamic vs static dimensions here?
    return tensor.bounding_shape(out_type=tf.int32)
  else:
    return tf.shape(tensor)
Returns the bounding shape of a Tensor, which is possibly Ragged.
language/canine/tensor_contracts.py