def register_ptransform(operation_def_subclass, tags=None):
  """Decorator to register a PTransform as the implementation for an analyzer.

  This function is used to define implementations of the analyzers defined in
  tensorflow_transform/analyzer_nodes.py and also the internal operations
  defined in tensorflow_transform/beam/beam_nodes.py.  The registered PTransform
  will be invoked as follows:

  outputs = inputs | operation.label >> MyPTransform(operation, extra_args)

  where operation is a the instance of the subclass that was registered,
  extra_args are global arguments available to each PTransform (see
  ConstructBeamPipelineVisitor.extra_args) and `inputs` is a tuple of
  PCollections corresponding to the inputs of the OperationNode being
  implemented.  The return value `outputs` should be a a tuple of PCollections
  corresponding to the outputs of the OperationNode.  If the OperationNode has
  a single output then the return value can also be a PCollection instead of a
  tuple.

  In some cases the implementation cannot be a PTransform and so instead the
  value being registered may also be a function.  The registered function will
  be invoked as follows:

  outputs = my_function(inputs, operation, extra_args)

  where inputs, operation, extra_args and outputs are the same as for the
  PTransform case.

  Args:
    operation_def_subclass: The class of attributes that is being registered.
        Should be a subclass of `tensorflow_transform.nodes.OperationDef`.
    tags: A set of string tags belonging to `EnvironmentTags`. If
        provided, the PTransform will be registered against all of them.

  Returns:
    A class decorator that registers a PTransform or function as an
        implementation of the OperationDef subclass.
  """

  def register(ptransform_class):
    assert isinstance(ptransform_class, type)
    assert issubclass(ptransform_class, beam.PTransform)
    assert tags is None or (tag in _ALLOWED_PTRANSFORM_TAGS for tag in tags)
    _PTRANSFORM_BY_OPERATION_DEF_SUBCLASS[
        operation_def_subclass].add_ptransform(ptransform_class, tags)
    return ptransform_class

  return register
Decorator to register a PTransform as the implementation for an analyzer.

  This function is used to define implementations of the analyzers defined in
  tensorflow_transform/analyzer_nodes.py and also the internal operations
  defined in tensorflow_transform/beam/beam_nodes.py.  The registered PTransform
  will be invoked as follows:

  outputs = inputs | operation.label >> MyPTransform(operation, extra_args)

  where operation is a the instance of the subclass that was registered,
  extra_args are global arguments available to each PTransform (see
  ConstructBeamPipelineVisitor.extra_args) and `inputs` is a tuple of
  PCollections corresponding to the inputs of the OperationNode being
  implemented.  The return value `outputs` should be a a tuple of PCollections
  corresponding to the outputs of the OperationNode.  If the OperationNode has
  a single output then the return value can also be a PCollection instead of a
  tuple.

  In some cases the implementation cannot be a PTransform and so instead the
  value being registered may also be a function.  The registered function will
  be invoked as follows:

  outputs = my_function(inputs, operation, extra_args)

  where inputs, operation, extra_args and outputs are the same as for the
  PTransform case.

  Args:
    operation_def_subclass: The class of attributes that is being registered.
        Should be a subclass of `tensorflow_transform.nodes.OperationDef`.
    tags: A set of string tags belonging to `EnvironmentTags`. If
        provided, the PTransform will be registered against all of them.

  Returns:
    A class decorator that registers a PTransform or function as an
        implementation of the OperationDef subclass.
  
tensorflow_transform/beam/common.py