Tags dataset-composition with extension point bug.
@@ -103,6 +103,7 @@ def compose_dataset_computation_with_computation(
raise TypeError(
'Expected a `tff.SequenceType` to be returned from '
'`dataset_computation`; found {} instead.'.format(dataset_return_type))
+ # TODO(b/226637447): This restriction seems unnecessary, and can be removed.
if dataset_computation.type_signature.parameter is None:
raise TypeError('Can only construct a new iterative process if '
'`dataset_computation` accepts a non-None arg; the '
@@ -309,6 +310,7 @@ def compose_dataset_computation_with_iterative_process(
raise TypeError(
'Expected a `tff.SequenceType` to be returned from '
'`dataset_computation`; found {} instead.'.format(dataset_return_type))
+ # TODO(b/226637447): This restriction seems unnecessary, and can be removed.
if dataset_computation.type_signature.parameter is None:
raise TypeError('Can only construct a new iterative process if '
'`dataset_computation` accepts a non-None arg; the '
