code: def make_transform(
      observation_transform: types.TensorTransformation,
      transform_next_observation: bool = True,
  ) -> reverb_dataset.Transform:
      """Creates the appropriate dataset transform for the given signature."""
  
      if transform_next_observation:
  
          def transform(x: reverb.ReplaySample) -> reverb.ReplaySample:
              return x._replace(
                  data=x.data._replace(
                      observation=observation_transform(
                          x.data.observation
                      ),
                      next_observation=observation_transform(
                          x.data.next_observation
                      ),
                  )
              )
  
      else:
  
          def transform(x: reverb.ReplaySample) -> reverb.ReplaySample:
              return x._replace(
                  data=x.data._replace(
                      observation=observation_transform(
                          x.data.observation
                      )
                  )
              )
  
      return transform
  
file_path: acme/acme/datasets/image_augmentation.py
imports: import enum
  from acme import types
  from acme.datasets import reverb
  import reverb
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: make_transform(observation_transform, transform_next_observation) and returns: transform, x, x