def METHOD_NAME(
    transformation: types.TensorValuedCallable
    ) -> snt.Module:
  """Convert a tensor transformation to a Sonnet Module.
  Args:
    transformation: A Callable that takes one or more (nested) Tensors, and
      returns one or more (nested) Tensors.
  Returns:
    A Sonnet Module that wraps the transformation.
  """
  if isinstance(transformation, snt.Module):
    return transformation
  module = TransformationWrapper(transformation)
  # Wrap the module to allow it to return an empty variable tuple.
  return snt.allow_empty_variables(module)