return builder.create<TF::SpaceToDepthOp>(
      cluster_func.getLoc(), transform_result_type, input, block_size);
}

// Performs transformation for a non-replicated input.
TF::SpaceToDepthOp HandleHostInput(Value input, int64_t index,
                                   tf_device::ClusterFuncOp cluster_func,
                                   int32_t block_size,
                                   ArrayRef<int64_t> input_shape) {