def _get_component_tensors(
    tensor: Union[tf.Tensor, composite_tensor.CompositeTensor]
) -> Iterable[tf.Tensor]:
  """Get all component tensors.

  Args:
    tensor: A `Tensor` or `CompositeTensor`.

  Returns:
    All `Tensor` components of `tensor`.

  Raises:
    ValueError if supplied `tensor` parameter is neither a `Tensor` nor a
    `CompositeTensor`.
  """
  if isinstance(tensor, tf.Tensor):
    return [tensor]
  elif isinstance(tensor, composite_tensor.CompositeTensor):
    return tf.nest.flatten(tensor, expand_composites=True)
  else:
    raise ValueError(
        'Unsupported tensor. Arg `tensor` is neither a `Tensor` nor a '
        f'`CompositeTensor`: {tensor}.')
Get all component tensors.

  Args:
    tensor: A `Tensor` or `CompositeTensor`.

  Returns:
    All `Tensor` components of `tensor`.

  Raises:
    ValueError if supplied `tensor` parameter is neither a `Tensor` nor a
    `CompositeTensor`.
  
tensorflow_transform/saved/saved_transform_io_v2.py