jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _fake_quant_feature(self, feature, scope, op, type, clamp_value=None):
        if clamp_value is None:
            clamp_value = self._clamp_value
        
        with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
            with self._graph.gradient_override_map({'Round': 'Identity'}):
                feature_abs_max = tf.reduce_max(tf.abs(feature, name=_MNN_mark_start_name))
                if _MNN_QAT_DEBUG:
                    print_op = tf.print("op:", op.name, type, "feature_abs_max, clamp_value:", feature_abs_max, clamp_value)
                    with tf.control_dependencies([print_op]):
                        instance_scale = feature_abs_max / clamp_value
                else:
                    instance_scale = feature_abs_max / clamp_value
                moving_average_scale = tf.get_variable(name='moving_average_scale', initializer=0., trainable=True)
                tf.add_to_collection(_MNN_variable_collection_name, moving_average_scale)

                def init():
                    init_assign = tf.assign(moving_average_scale, instance_scale)
                    initialize_op = tf.assign(self._op_scale_initialized[op][type], 1)
                    if _MNN_QAT_DEBUG:
                        print_op = tf.print("init", moving_average_scale.name, init_assign)
                        with tf.control_dependencies([print_op, initialize_op]):
                            return tf.identity(init_assign)
                    else:
                        with tf.control_dependencies([initialize_op]):
                            return tf.identity(init_assign)

                def do_nothing():
                    return moving_average_scale

                init_moving_scale = tf.cond(tf.equal(self._op_scale_initialized[op][type], 0), init, do_nothing)
                self._feature_scale_init_op[moving_average_scale] = init_moving_scale

                def fake_quant():
                    with tf.control_dependencies([init_moving_scale]):
                        s_scale = custom_scale_grad(moving_average_scale, tf.size(feature), clamp_value)
                        quant_x = tf.clip_by_value(tf.round(feature / s_scale), -clamp_value, clamp_value)
                        fake_quant_x = quant_x * s_scale
                        return fake_quant_x

                new_feature = tf.cond(tf.constant(True), fake_quant, lambda:feature, name=_MNN_mark_end_name)
        
        self._feature_scale_original_tensor_dims_map[moving_average_scale] = [feature, feature.shape.as_list()[-1]]
        
        return new_feature, moving_average_scale, feature.shape.as_list()[-1], init_moving_scale
//  [python]
def evaluate(model, tf_dataset, max_step=1000, num_eval_samples=None):
  """Evaluate the model."""
  if FLAGS.grad_type in ['vimco', 'local-disarm']:
    num_samples = FLAGS.num_pairs * 2
  elif num_eval_samples:
    num_samples = num_eval_samples
  elif FLAGS.num_eval_samples:
    num_samples = FLAGS.num_eval_samples
  else:
    num_samples = FLAGS.num_pairs
  tf.print('Evaluate with samples: ', num_samples)
  loss = 0.
  n = 0.
  for batch in tf_dataset.map(process_batch_input):
    if n >= max_step:  # used for train_ds, which is a `repeat` dataset.
      break
    if num_samples > 1:
      batch_size = tf.shape(batch)[0]
      input_batch = tf.tile(batch, [num_samples, 1])
      elbo = tf.reshape(model(input_batch)[0], [num_samples, batch_size])
      objectives = (tf.reduce_logsumexp(elbo, axis=0, keepdims=False) -
                    tf.math.log(tf.cast(tf.shape(elbo)[0], tf.float32)))
    else:
      objectives = model(batch)[0]
    loss -= tf.reduce_mean(objectives)
    n += 1.
  return loss / n
//  [python]
def log_validation_step(opts, file_writer, image, mask, step, pred, val_loss, val_miou, val_auc,epoch):
    """ Log to file writer after a validation step """
    if hvd.local_rank() == 0:
        with file_writer.as_default():
            tf.summary.image(f'Validation image of worker {hvd.rank()}', image, step=tf.cast(step, tf.int64),
                             max_outputs=5)
            tf.summary.image(f'Validation mask of worker {hvd.rank()}', mask, step=tf.cast(step, tf.int64),
                             max_outputs=5)
            tf.summary.image(f'Validation_prediction of worker {hvd.rank()}', pred, step=tf.cast(step, tf.int64),
                             max_outputs=2)
            tf.summary.scalar(f'Validation Loss of worker {hvd.rank()}', val_loss, step=tf.cast(step, tf.int64))
            tf.summary.scalar(f'Validation Mean IoU of worker {hvd.rank()}', val_miou.result().numpy(), step=tf.cast(step, tf.int64))
            tf.summary.scalar('Validation AUC', val_auc.result().numpy(), step=tf.cast(step, tf.int64))

        file_writer.flush()

        tf.compat.v1.logging.info('Validation at epoch %d Step %d / %d : validation loss %f : validation miou %f : validation auc %f', epoch, step, opts.steps_per_epoch, val_loss, val_miou.result().numpy(), val_auc.result().numpy())
    return
//  [python]
def log_validation_step(opts, file_writer, image, mask, step, pred, val_loss, val_miou, val_auc,epoch):
    """ Log to file writer after a validation step """
    if hvd.local_rank() == 0:
        with file_writer.as_default():
            tf.summary.image(f'Validation image of worker {hvd.rank()}', image, step=tf.cast(step, tf.int64),
                             max_outputs=5)
            tf.summary.image(f'Validation mask of worker {hvd.rank()}', mask, step=tf.cast(step, tf.int64),
                             max_outputs=5)
            tf.summary.image(f'Validation_prediction of worker {hvd.rank()}', pred, step=tf.cast(step, tf.int64),
                             max_outputs=2)
            tf.summary.scalar(f'Validation Loss of worker {hvd.rank()}', val_loss, step=tf.cast(step, tf.int64))
            tf.summary.scalar(f'Validation Mean IoU of worker {hvd.rank()}', val_miou.result().numpy(), step=tf.cast(step, tf.int64))
            tf.summary.scalar('Validation AUC', val_auc.result().numpy(), step=tf.cast(step, tf.int64))

        file_writer.flush()

        tf.print('Validation at epoch',epoch,'Step', step, '/', opts.steps_per_epoch,
                 ': validation loss', val_loss,
                 ': validation miou', val_miou.result().numpy(),
                 ': validation auc', val_auc.result().numpy())
    return
//  [python]
def log_training_step(opts, model, file_writer, x, y, loss, pred, step, metrics, optimizer, steptime,epoch):
    """ Log to file writer during training"""
    if hvd.local_rank() == 0 and hvd.rank() == 0:

        # Make y [batch_size,image_size,image_size,1], prepare for metrics
        y = tf.argmax(y,axis=-1)[...,None]
        compute_loss, compute_miou, compute_auc = metrics
        
        compute_miou.update_state(y,pred)
        compute_auc.update_state(y,pred)

        # Training Prints
        tf.compat.v1.logging.info('\nEpoch:',epoch,'Step', step, '/', opts.steps_per_epoch,
                 ': loss', loss,
                 ': miou', compute_miou.result().numpy(),
                 ': auc', compute_auc.result().numpy(), '\n')
        
        with file_writer.as_default():

            image = tf.cast(255 * x, tf.uint8)
            mask = tf.cast(255 * y, tf.uint8)
            summary_predictions = tf.cast(tf.expand_dims(pred * 255, axis=-1), tf.uint8)
            tf.summary.scalar('Training StepTime', steptime, step=tf.cast(step, tf.int64))
            tf.summary.image('Train_image', image, step=tf.cast(step, tf.int64), max_outputs=2)
            tf.summary.image('Train_mask', mask, step=tf.cast(step, tf.int64), max_outputs=2)
            tf.summary.image('Train_prediction', summary_predictions, step=tf.cast(step, tf.int64),
                             max_outputs=2)
            tf.summary.scalar('Training Loss', loss, step=tf.cast(step, tf.int64))
            tf.summary.scalar('Training mIoU', compute_miou.result().numpy(),step=tf.cast(step, tf.int64))
            tf.summary.scalar('Training AUC', compute_auc.result().numpy(), step=tf.cast(step, tf.int64))

            # Logging the optimizer's hyperparameters
            for key in optimizer._hyper:
                tf.summary.scalar(key, optimizer._hyper[key].numpy(), step=tf.cast(step, tf.int64))
            # Extract weights and filter out None elemens for aspp without weights
            weights = filter(None, [x.weights for x in model.layers])
            for var in weights:
                tf.summary.histogram('%s' % var[0].name, var[0], step=tf.cast(step, tf.int64))

        file_writer.flush()

    return
//  [python]
def log_training_step(opts, model, file_writer, x, y, loss, pred, step, metrics, optimizer, steptime,epoch):
    """ Log to file writer during training"""
    if hvd.local_rank() == 0 and hvd.rank() == 0:

        # Make y [batch_size,image_size,image_size,1], prepare for metrics
        y = tf.argmax(y,axis=-1)[...,None]
        compute_loss, compute_miou, compute_auc = metrics
        
        compute_miou.update_state(y,pred)
        compute_auc.update_state(y,pred)

        # Training Prints
        tf.print('\nEpoch:',epoch,'Step', step, '/', opts.steps_per_epoch,
                 ': loss', loss,
                 ': miou', compute_miou.result().numpy(),
                 ': auc', compute_auc.result().numpy(), '\n')
        
        with file_writer.as_default():

            image = tf.cast(255 * x, tf.uint8)
            mask = tf.cast(255 * y, tf.uint8)
            summary_predictions = tf.cast(tf.expand_dims(pred * 255, axis=-1), tf.uint8)
            tf.summary.scalar('Training StepTime', steptime, step=tf.cast(step, tf.int64))
            tf.summary.image('Train_image', image, step=tf.cast(step, tf.int64), max_outputs=2)
            tf.summary.image('Train_mask', mask, step=tf.cast(step, tf.int64), max_outputs=2)
            tf.summary.image('Train_prediction', summary_predictions, step=tf.cast(step, tf.int64),
                             max_outputs=2)
            tf.summary.scalar('Training Loss', loss, step=tf.cast(step, tf.int64))
            tf.summary.scalar('Training mIoU', compute_miou.result().numpy(),step=tf.cast(step, tf.int64))
            tf.summary.scalar('Training AUC', compute_auc.result().numpy(), step=tf.cast(step, tf.int64))

            # Logging the optimizer's hyperparameters
            for key in optimizer._hyper:
                tf.summary.scalar(key, optimizer._hyper[key].numpy(), step=tf.cast(step, tf.int64))
            # Extract weights and filter out None elemens for aspp without weights
            weights = filter(None, [x.weights for x in model.layers])
            for var in weights:
                tf.summary.histogram('%s' % var[0].name, var[0], step=tf.cast(step, tf.int64))

        file_writer.flush()

    return
//  [python]
def building_block(self, xyz, feature, neigh_idx, d_out, name, is_training):
        d_in = feature.get_shape()[-1]
        f_xyz = self.relative_pos_encoding(xyz, neigh_idx)
        print_op = tf.Print(f_xyz, [tf.shape(input=f_xyz)])
        f_xyz = helper_tf_util.conv2d(f_xyz, d_in, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training)
        f_neighbours = self.gather_neighbour(tf.squeeze(feature, axis=2), neigh_idx)
        f_concat = tf.concat([f_neighbours, f_xyz], axis=-1)
        f_pc_agg = self.att_pooling(f_concat, d_out // 2, name + 'att_pooling_1', is_training)

        f_xyz = helper_tf_util.conv2d(f_xyz, d_out // 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training)
        f_neighbours = self.gather_neighbour(tf.squeeze(f_pc_agg, axis=2), neigh_idx)
        f_concat = tf.concat([f_neighbours, f_xyz], axis=-1)
        f_pc_agg = self.att_pooling(f_concat, d_out, name + 'att_pooling_2', is_training)
        return f_pc_agg, print_op
//  [python]
def building_block(self, xyz, feature, neigh_idx, d_out, name, is_training):
        d_in = feature.get_shape()[-1]
        f_xyz = self.relative_pos_encoding(xyz, neigh_idx)
        print_op = tf.print(tf.shape(input=f_xyz),output_stream=sys.stderr)
        f_xyz = helper_tf_util.conv2d(f_xyz, d_in, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training)
        f_neighbours = self.gather_neighbour(tf.squeeze(feature, axis=2), neigh_idx)
        f_concat = tf.concat([f_neighbours, f_xyz], axis=-1)
        f_pc_agg = self.att_pooling(f_concat, d_out // 2, name + 'att_pooling_1', is_training)

        f_xyz = helper_tf_util.conv2d(f_xyz, d_out // 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training)
        f_neighbours = self.gather_neighbour(tf.squeeze(f_pc_agg, axis=2), neigh_idx)
        f_concat = tf.concat([f_neighbours, f_xyz], axis=-1)
        f_pc_agg = self.att_pooling(f_concat, d_out, name + 'att_pooling_2', is_training)
        return f_pc_agg, print_op
//  [python]
def train():    
    # 使用TensorFlow内置的tf.data预处理数据集
    (x, y), _ = tf.keras.datasets.mnist.load_data()
    dataset = tf.data.Dataset.from_tensor_slices((x, y))
    dataset = dataset.map(prepare_mnist_features_and_labels)
    dataset = dataset.take(20000).shuffle(20000).batch(batch_size)
    batch_index = 0
    for X, y in dataset:        
        with tf.GradientTape() as tape:
            y_pred = model(X)
            loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
            loss = tf.reduce_mean(loss)
        tf.compat.v1.print("batch", batch_index, ", loss", loss)
        grads = tape.gradient(loss, model.variables)
        optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
        batch_index += 1
//  [python]
def train():    
    # 使用TensorFlow内置的tf.data预处理数据集
    (x, y), _ = tf.keras.datasets.mnist.load_data()
    dataset = tf.data.Dataset.from_tensor_slices((x, y))
    dataset = dataset.map(prepare_mnist_features_and_labels)
    dataset = dataset.take(20000).shuffle(20000).batch(batch_size)
    batch_index = 0
    for X, y in dataset:        
        with tf.GradientTape() as tape:
            y_pred = model(X)
            loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
            loss = tf.reduce_mean(loss)
        tf.print("batch", batch_index, ", loss", loss)
        grads = tape.gradient(loss, model.variables)
        optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
        batch_index += 1
//  [python]
def call(self, inputs, **kwargs):
        tf.print(self.log_id, inputs)
        return inputs
//  [python]
def test_dicom_image_concurrency():
    """test_decode_dicom_image_currency"""

    @tf.function
    def preprocess(dcm_content):
        tags = tfio.image.decode_dicom_data(
            dcm_content, tags=[tfio.image.dicom_tags.PatientsName]
        )
        tf.print(tags)
        image = tfio.image.decode_dicom_image(dcm_content, dtype=tf.float32)
        return image

    dcm_path = os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "test_dicom",
        "TOSHIBA_J2K_OpenJPEGv2Regression.dcm",
    )

    dataset = (
        tf.data.Dataset.from_tensor_slices([dcm_path])
        .repeat()
        .map(tf.io.read_file)
        .map(preprocess, num_parallel_calls=8)
        .take(200)
    )
    for i, item in enumerate(dataset):
        print(tf.shape(item), i)
        assert np.array_equal(tf.shape(item), [1, 512, 512, 1])

    dcm_path = os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "test_dicom",
        "US-PAL-8-10x-echo.dcm",
    )

    dataset = (
        tf.data.Dataset.from_tensor_slices([dcm_path])
        .repeat()
        .map(tf.io.read_file)
        .map(preprocess, num_parallel_calls=8)
        .take(200)
    )
    for i, item in enumerate(dataset):
        print(tf.shape(item), i)
        assert np.array_equal(tf.shape(item), [10, 430, 600, 3])
//  [python]
mirrored_strategy = tf.distribute.MirroredStrategy()
BATCH_SIZE = params['batch_size'] * mirrored_strategy.num_replicas_in_sync

train_ds = train.to_tf_dataset().map(parse_labeled_data).shuffle(BATCH_SIZE*100).batch(BATCH_SIZE).repeat()
val_ds = val.to_tf_dataset().map(parse_labeled_data).batch(BATCH_SIZE)

for record in train_ds:
  tf.print(record)
  break
//  [python]
def building_block(self, xyz, feature, neigh_idx, d_out, name, is_training):
        d_in = feature.get_shape()[-1].value
        f_xyz = self.relative_pos_encoding(xyz, neigh_idx)
        print_op = tf.Print(f_xyz, [tf.shape(f_xyz)])
        f_xyz = helper_tf_util.conv2d(f_xyz, d_in, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training)
        f_neighbours = self.gather_neighbour(tf.squeeze(feature, axis=2), neigh_idx)
        f_concat = tf.concat([f_neighbours, f_xyz], axis=-1)
        f_pc_agg = self.att_pooling(f_concat, d_out // 2, name + 'att_pooling_1', is_training)

        f_xyz = helper_tf_util.conv2d(f_xyz, d_out // 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training)
        f_neighbours = self.gather_neighbour(tf.squeeze(f_pc_agg, axis=2), neigh_idx)
        f_concat = tf.concat([f_neighbours, f_xyz], axis=-1)
        f_pc_agg = self.att_pooling(f_concat, d_out, name + 'att_pooling_2', is_training)
        return f_pc_agg, print_op
//  [python]
def building_block(self, xyz, feature, neigh_idx, d_out, name, is_training):
        d_in = feature.get_shape()[-1].value
        f_xyz = self.relative_pos_encoding(xyz, neigh_idx)
        print_op = tf.print(tf.shape(f_xyz),output_stream=sys.stderr)
        f_xyz = helper_tf_util.conv2d(f_xyz, d_in, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training)
        f_neighbours = self.gather_neighbour(tf.squeeze(feature, axis=2), neigh_idx)
        f_concat = tf.concat([f_neighbours, f_xyz], axis=-1)
        f_pc_agg = self.att_pooling(f_concat, d_out // 2, name + 'att_pooling_1', is_training)

        f_xyz = helper_tf_util.conv2d(f_xyz, d_out // 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training)
        f_neighbours = self.gather_neighbour(tf.squeeze(f_pc_agg, axis=2), neigh_idx)
        f_concat = tf.concat([f_neighbours, f_xyz], axis=-1)
        f_pc_agg = self.att_pooling(f_concat, d_out, name + 'att_pooling_2', is_training)
        return f_pc_agg, print_op
//  [python]
def recovery_to_raw_image_shape(path_to_raw_img, path_to_produce_img):
    raw_img = tf.io.read_file(path_to_raw_img)
    raw_img = tf.image.decode_image(raw_img, channels=3)
    tf.print(f"raw_img_shape {tf.shape(raw_img)[:-1]}")

    produce_img = tf.io.read_file(path_to_produce_img)
    produce_img = tf.image.decode_image(produce_img, channels=3)
    produce_img = tf.image.convert_image_dtype(produce_img, tf.float32)
    tf.print(f"produce_img_shape {tf.shape(produce_img)[:-1]}")
    produce_img = tf.image.resize(produce_img, tf.shape(raw_img)[:-1])
    tf.print(f"recovery produce_img_shape {tf.shape(produce_img)[:-1]}")
    produce_img = produce_img[tf.newaxis, :]
    return produce_img
//  [python]
def train_step(image_data, target):
    with tf.GradientTape() as tape:
        pred_result = model(image_data, training=True)
        giou_loss=conf_loss=prob_loss=0

        # optimizing process
        for i in range(3):
            conv, pred = pred_result[i*2], pred_result[i*2+1]
            loss_items = compute_loss(pred, conv, *target[i], i)
            giou_loss += loss_items[0]
            conf_loss += loss_items[1]
            prob_loss += loss_items[2]

        total_loss = giou_loss + conf_loss + prob_loss

        gradients = tape.gradient(total_loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
        tf.print("=> STEP %4d   lr: %.6f   giou_loss: %4.2f   conf_loss: %4.2f   "
                 "prob_loss: %4.2f   total_loss: %4.2f" %(global_steps, optimizer.lr.numpy(),
                                                          giou_loss, conf_loss,
                                                          prob_loss, total_loss))
        # update learning rate
        global_steps.assign_add(1)
        if global_steps < warmup_steps:
            lr = global_steps / warmup_steps *cfg.TRAIN.LR_INIT
        else:
            lr = cfg.TRAIN.LR_END + 0.5 * (cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * (
                (1 + tf.cos((global_steps - warmup_steps) / (total_steps - warmup_steps) * np.pi))
            )
        optimizer.lr.assign(lr.numpy())

        # writing summary data
        with writer.as_default():
            tf.summary.scalar("lr", optimizer.lr, step=global_steps)
            tf.summary.scalar("loss/total_loss", total_loss, step=global_steps)
            tf.summary.scalar("loss/giou_loss", giou_loss, step=global_steps)
            tf.summary.scalar("loss/conf_loss", conf_loss, step=global_steps)
            tf.summary.scalar("loss/prob_loss", prob_loss, step=global_steps)
        writer.flush()
//  [python]
config = tfe.get_config()

if len(sys.argv) > 1:

    #
    # assume we're running as a server
    #

    player_name = str(sys.argv[1])

    server = config.server(player_name)
    server.start()
    server.join()

else:

    #
    # assume we're running as master
    #

    def provide_weights() -> tf.Tensor:
        raw_w = np.array([5, 5, 5, 5]).reshape((2, 2))
        w = tf.constant(raw_w)
        t.print(w, [w])
        return w

    def provide_input() -> tf.Tensor:
        x = tf.constant([1, 2, 3, 4], shape=(2, 2), dtype=tf.float32)
        t.print(x, [x])
        return x

    def receive_output(prediction):

        t.print([], [prediction], summarize=4)
        return []

    with tfe.protocol.Pond() as prot:

        # treat weights as private
        w = prot.define_private_input("model-provider", provide_weights)

        # load input for prediction
        x = prot.define_private_input("input-provider", provide_input)

        # compute prediction
        y = prot.matmul(x, w)

        # send output
        prediction_op = prot.define_output("input-provider", y, receive_output)

        with tfe.Session() as sess:
            sess.run(tf.global_variables_initializer(), tag="init")

            for _ in range(5):
                sess.run(prediction_op, tag="prediction")
//  [python]
config = tfe.get_config()

if len(sys.argv) > 1:

    #
    # assume we're running as a server
    #

    player_name = str(sys.argv[1])

    server = config.server(player_name)
    server.start()
    server.join()

else:

    #
    # assume we're running as master
    #

    def provide_weights() -> tf.Tensor:
        raw_w = np.array([5, 5, 5, 5]).reshape((2, 2))
        w = tf.constant(raw_w)
        tf.print(w, [w])
        return w

    def provide_input() -> tf.Tensor:
        x = tf.constant([1, 2, 3, 4], shape=(2, 2), dtype=tf.float32)
        tf.print(x, [x])
        return x

    def receive_output(prediction):

        tf.print([], [prediction], summarize=4)
        return []

    with tfe.protocol.Pond() as prot:

        # treat weights as private
        w = prot.define_private_input("model-provider", provide_weights)

        # load input for prediction
        x = prot.define_private_input("input-provider", provide_input)

        # compute prediction
        y = prot.matmul(x, w)

        # send output
        prediction_op = prot.define_output("input-provider", y, receive_output)

        with tfe.Session() as sess:
            sess.run(tf.global_variables_initializer(), tag="init")

            for _ in range(5):
                sess.run(prediction_op, tag="prediction")
//  [python]
def print_one_line(*args):
    t.print("\033[K", end="")
    t.print("\r", *args, sep="", end=" ", output_stream=sys.stdout)
//  [python]
def print_one_line(*args):
    tf.print("\033[K", end="")
    tf.print("\r", *args, sep="", end=" ", output_stream=sys.stdout)
//  [python]
def train():    
    # 使用TensorFlow內建的tf.data預處理資料集
    (x, y), _ = tf.keras.datasets.mnist.load_data()
    dataset = tf.data.Dataset.from_tensor_slices((x, y))
    dataset = dataset.map(prepare_mnist_features_and_labels)
    dataset = dataset.take(20000).shuffle(20000).batch(batch_size)
    batch_index = 0
    for X, y in dataset:        
        with tf.GradientTape() as tape:
            y_pred = model(X)
            loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
            loss = tf.reduce_mean(loss)
        tf.compat.v1.logging.info("batch %d, loss %f", batch_index, loss)
        grads = tape.gradient(loss, model.variables)
        optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
        batch_index += 1
//  [python]
def train():    
    # 使用TensorFlow內建的tf.data預處理資料集
    (x, y), _ = tf.keras.datasets.mnist.load_data()
    dataset = tf.data.Dataset.from_tensor_slices((x, y))
    dataset = dataset.map(prepare_mnist_features_and_labels)
    dataset = dataset.take(20000).shuffle(20000).batch(batch_size)
    batch_index = 0
    for X, y in dataset:        
        with tf.GradientTape() as tape:
            y_pred = model(X)
            loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
            loss = tf.reduce_mean(loss)
        tf.print("batch", batch_index, ", loss", loss)
        grads = tape.gradient(loss, model.variables)
        optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
        batch_index += 1
//  [python]
def print_metrics(metrics, step, every, name='metrics'):
  means, updates = [], []
  for key, value in metrics.items():
    key = 'metrics_{}_{}'.format(name, key)
    mean = tools.StreamingMean((), tf.float32, key)
    means.append(mean)
    updates.append(mean.submit(value))
  with tf.control_dependencies(updates):
    # message = 'step/' + '/'.join(metrics.keys()) + ' = '
    message = '{}: step/{} ='.format(name, '/'.join(metrics.keys()))
    gs = tf.train.get_or_create_global_step()
    print_metrics = tf.cond(
        tf.equal(step % every, 0),
        lambda: tf.Print(message, [gs] + [mean.clear() for mean in means]),
        tf.no_op)
  return print_metrics
//  [python]
def print_metrics(metrics, step, every, name='metrics'):
  means, updates = [], []
  for key, value in metrics.items():
    key = 'metrics_{}_{}'.format(name, key)
    mean = tools.StreamingMean((), tf.float32, key)
    means.append(mean)
    updates.append(mean.submit(value))
  with tf.control_dependencies(updates):
    # message = 'step/' + '/'.join(metrics.keys()) + ' = '
    message = '{}: step/{} ='.format(name, '/'.join(metrics.keys()))
    gs = tf.train.get_or_create_global_step()
    print_metrics = tf.cond(
        tf.equal(step % every, 0),
        lambda: tf.print(message, [gs] + [mean.clear() for mean in means]),
        tf.no_op)
  return print_metrics
//  [python]
def multiple_prints(x, y):
  t.print('x is', x)
  t.print('y is', y)
//  [python]
def multiple_prints(x, y):
  tf.print('x is', x)
  tf.print('y is', y)
//  [python]
def tf_print(x):
  t.print(x)
//  [python]
def tf_print(x):
  tf.print(x)
//  [python]
def test_b279454591(self):
    """Test case when tensorflow function returns `StatefulPartitionedCall` op."""
    inputs = jnp.ones([10], dtype=jnp.float32)

    # With one or more outputs, it is okay.
    def tf_f(x):
      y = tf.math.sin(3.0)
      tf.print(y)
      return x

    jax_f = jax2tf.call_tf(
        tf.function(tf_f),
        call_tf_graph=True,
    )
    tf_f_rt = jax2tf.convert(
        jax_f,
        native_serialization=True,
        with_gradient=False,
    )
    _, _ = tf_test_util.SaveAndLoadFunction(tf_f_rt, input_args=[inputs])

    # With zero output, it return `StatefulPartitionedCall` op instead.
    def tf_f_2():
      y = tf.math.sin(3.0)
      tf.print(y)
      return

    jax_f_2 = jax2tf.call_tf(tf.function(tf_f_2), call_tf_graph=True)
    tf_f_rt_2 = jax2tf.convert(
        jax_f_2,
        native_serialization=True,
        with_gradient=False,
    )
    _, _ = tf_test_util.SaveAndLoadFunction(tf_f_rt_2, input_args=[])
//  [python]
def _fake_quant_weight(self, weight_tensor, reduce_dims, scope, clamp_value=None):
        if clamp_value is None:
            clamp_value = self._clamp_value

        with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
            with self._graph.gradient_override_map({'Round': 'Identity'}):
                scales = tf.reduce_max(tf.abs(weight_tensor, name=_MNN_mark_start_name), axis=reduce_dims, keep_dims=True) / clamp_value + self._eps
                weight_scales = tf.get_variable(name='weight_scales', initializer=tf.zeros_like(scales), trainable=False)
                tf.add_to_collection(_MNN_variable_collection_name, weight_scales)
                scales_assign = tf.assign(weight_scales, scales)
                with tf.control_dependencies([scales_assign]):
                    quant_w = tf.clip_by_value(tf.round(weight_tensor / scales), -clamp_value, clamp_value)
                    fake_quant_w = quant_w * scales

            if _MNN_QAT_DEBUG:
                print_op = tf.Print(scope, [scope, "weight clamp value:", clamp_value])
                with tf.control_dependencies([print_op]):
                    new_weight_tensor = tf.cond(tf.constant(True), lambda: fake_quant_w, lambda: weight_tensor, name=_MNN_mark_end_name)
            else:
                new_weight_tensor = tf.cond(tf.constant(True), lambda: fake_quant_w, lambda: weight_tensor, name=_MNN_mark_end_name)

        self._weight_scale_original_tensor_map[weight_scales] = weight_tensor

        return new_weight_tensor, weight_scales
//  [python]
def _fake_quant_weight(self, weight_tensor, reduce_dims, scope, clamp_value=None):
        if clamp_value is None:
            clamp_value = self._clamp_value

        with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
            with self._graph.gradient_override_map({'Round': 'Identity'}):
                scales = tf.reduce_max(tf.abs(weight_tensor, name=_MNN_mark_start_name), axis=reduce_dims, keep_dims=True) / clamp_value + self._eps
                weight_scales = tf.get_variable(name='weight_scales', initializer=tf.zeros_like(scales), trainable=False)
                tf.add_to_collection(_MNN_variable_collection_name, weight_scales)
                scales_assign = tf.assign(weight_scales, scales)
                with tf.control_dependencies([scales_assign]):
                    quant_w = tf.clip_by_value(tf.round(weight_tensor / scales), -clamp_value, clamp_value)
                    fake_quant_w = quant_w * scales

            if _MNN_QAT_DEBUG:
                print_op = tf.print(scope, "weight clamp value:", clamp_value)
                with tf.control_dependencies([print_op]):
                    new_weight_tensor = tf.cond(tf.constant(True), lambda: fake_quant_w, lambda: weight_tensor, name=_MNN_mark_end_name)
            else:
                new_weight_tensor = tf.cond(tf.constant(True), lambda: fake_quant_w, lambda: weight_tensor, name=_MNN_mark_end_name)

        self._weight_scale_original_tensor_map[weight_scales] = weight_tensor

        return new_weight_tensor, weight_scales
//  [python]
def f(x):
    print("The function is running in Python")
    tf.print(x)
//  [python]
def fit(
        self, train_datasets, epochs, c_round, rounds, num_batches=None, T=1, verbose=0
    ):
        """Trains the model on given datasets for a specified number of epochs.

        This custom training loop handles both labeled and unlabeled data, applying
        pseudo-labeling for semi-supervised learning.

        Parameters
        ----------
        - train_datasets (tuple): A tuple containing labeled
                                  and unlabeled training datasets.
        - epochs (int): Number of epochs to train the model.
        - c_round (int): Current round of training in a federated learning scenario.
        - rounds (int): Total number of rounds in the federated learning scenario.
        - num_batches (int, optional): Number of batches per epoch.
        - T (float, optional): Temperature parameter for softmax in pseudo-labeling.
        - verbose (int, optional): Verbosity mode.

        Returns
        -------
        - History: An object containing training history information.
        """
        l_dataset, u_dataset = train_datasets
        self.history.on_train_begin()
        for epoch in range(epochs):
            start, epoch_loss = time.time(), 0.0
            # Threshold value to get good psuedo labels, part of Eq-5
            self.confidence = __class__.cosine_confidence_scheduler(
                step=c_round, total_steps=rounds
            )
            progbar = tf.keras.utils.Progbar(
                target=num_batches - 1,
                width=30,
                verbose=verbose if verbose != 2 else 0,
                interval=0.05,
                stateful_metrics=None,
                unit_name="step",
            )
            if verbose == 1:
                print(f"Epoch {epoch+1}/{epochs}")
            step = 0
            # Run train step on each batch
            # tf.data.Dataset.repeat
            # (https://www.tensorflow.org/api_docs/python/tf/data/Dataset#repeat)
            # can also be used instead of itertools.cycle
            for step, ((l_batch_x, l_batch_y), (u_batch_x, u_batch_y)) in enumerate(
                zip(itertools.cycle(l_dataset), u_dataset)
            ):
                batch_x, batch_y = tf.concat([l_batch_x, u_batch_x], 0), tf.concat(
                    [l_batch_y, u_batch_y], 0
                )
                loss = self.train_step(x=batch_x, y=batch_y, T=T)
                progbar.update(current=step, values=[("loss", float(loss["loss"]))])
                epoch_loss += float(loss["loss"])
            # Store result in history
            self.history.on_epoch_end(
                epoch=epoch,
                logs={"loss": epoch_loss / num_batches, "accuracy": float(0.0)},
            )
            if verbose == 2:
                print(
                    f"Epoch {epoch + 1}/{epochs}\n"
                    f"{step}/{step} - {time.time() - start:.1f}s - "
                    f"loss: {epoch_loss / num_batches:.4f} - "
                    f"accuracy: {float(0.0):.4f}"
                )
        return self.history
//  [python]
def fit(
        self, train_datasets, epochs, c_round, rounds, num_batches=None, T=1, verbose=0
    ):
        """Trains the model on given datasets for a specified number of epochs.

        This custom training loop handles both labeled and unlabeled data, applying
        pseudo-labeling for semi-supervised learning.

        Parameters
        ----------
        - train_datasets (tuple): A tuple containing labeled
                                  and unlabeled training datasets.
        - epochs (int): Number of epochs to train the model.
        - c_round (int): Current round of training in a federated learning scenario.
        - rounds (int): Total number of rounds in the federated learning scenario.
        - num_batches (int, optional): Number of batches per epoch.
        - T (float, optional): Temperature parameter for softmax in pseudo-labeling.
        - verbose (int, optional): Verbosity mode.

        Returns
        -------
        - History: An object containing training history information.
        """
        l_dataset, u_dataset = train_datasets
        self.history.on_train_begin()
        for epoch in range(epochs):
            start, epoch_loss = time.time(), 0.0
            # Threshold value to get good psuedo labels, part of Eq-5
            self.confidence = __class__.cosine_confidence_scheduler(
                step=c_round, total_steps=rounds
            )
            progbar = tf.keras.utils.Progbar(
                target=num_batches - 1,
                width=30,
                verbose=verbose if verbose != 2 else 0,
                interval=0.05,
                stateful_metrics=None,
                unit_name="step",
            )
            if verbose == 1:
                tf.print(f"Epoch {epoch+1}/{epochs}")
            step = 0
            # Run train step on each batch
            # tf.data.Dataset.repeat
            # (https://www.tensorflow.org/api_docs/python/tf/data/Dataset#repeat)
            # can also be used instead of itertools.cycle
            for step, ((l_batch_x, l_batch_y), (u_batch_x, u_batch_y)) in enumerate(
                zip(itertools.cycle(l_dataset), u_dataset)
            ):
                batch_x, batch_y = tf.concat([l_batch_x, u_batch_x], 0), tf.concat(
                    [l_batch_y, u_batch_y], 0
                )
                loss = self.train_step(x=batch_x, y=batch_y, T=T)
                progbar.update(current=step, values=[("loss", float(loss["loss"]))])
                epoch_loss += float(loss["loss"])
            # Store result in history
            self.history.on_epoch_end(
                epoch=epoch,
                logs={"loss": epoch_loss / num_batches, "accuracy": float(0.0)},
            )
            if verbose == 2:
                tf.print(
                    f"Epoch {epoch + 1}/{epochs}\n"
                    f"{step}/{step} - {time.time() - start:.1f}s - "
                    f"loss: {epoch_loss / num_batches:.4f} - "
                    f"accuracy: {float(0.0):.4f}"
                )
        return self.history
//  [python]
def test_call_tf_graph_ordered(self, *, version: int):
    self.override_serialization_version(version)
    @tf.function
    def tf_print(x):
      tf.print(x)

    call_tf_print = jax2tf.call_tf(
        tf_print,
        call_tf_graph=True,
        ordered=True,
    )

    x = jnp.array(1.0, dtype=jnp.float32)

    def body(i, x):
      call_tf_print(x)
      return x + 1

    @jax.jit
    def f_jax(x):
      return jax.lax.fori_loop(0, 4, body, x)

    num_custom_calls = 0

    def _check_mlir_ops(op):
      nonlocal num_custom_calls

      if (
          op.operation.name == "stablehlo.custom_call"
          and ir.StringAttr(op.attributes["call_target_name"]).value
          == "tf.call_tf_function"
      ):
        num_custom_calls += 1

        # The custom call op must have `has_token_input_output` attribute.
        tf_backend_config = ir.DictAttr(op.attributes["tf.backend_config"])
        self.assertTrue(
            ir.BoolAttr(tf_backend_config["has_token_input_output"]).value
        )

        # Verify that the first argument/result of the custom call op is a token
        # type. This is a calling convention defined by `has_token_input_output`.
        self.assertTrue(hlo.TokenType.isinstance(op.operands[0].type))
        self.assertTrue(hlo.TokenType.isinstance(op.results[0].type))

    stablehlo_module = None
    with self.assertRaisesRegex(
        ValueError,
        "call_tf_graph=True only support exporting by jax2tf.convert currently",
    ):
      lower = f_jax.lower(x)
      self.assertNotEmpty(lower._lowering.compile_args["ordered_effects"])
      stablehlo_module = lower.compiler_ir("stablehlo")
    if stablehlo_module:
      self._walk_stablehlo_operations(stablehlo_module, _check_mlir_ops)
      self.assertEqual(num_custom_calls, 1)

    f_tf = jax2tf.convert(
        f_jax,
        native_serialization=True,
        with_gradient=False,
    )
    _, restored_model = tf_test_util.SaveAndLoadFunction(f_tf, input_args=[x])
//  [python]
def receive_output(tensor: tf.Tensor) -> tf.Tensor:
    tf.compat.v1.print(tensor, [tensor])
    return tensor
//  [python]
def receive_output(tensor: tf.Tensor) -> tf.Tensor:
    tf.print(tensor, [tensor])
    return tensor
//  [python]
def audio_to_features(
    audio,
    sample_rate,
    transcript=None,
    clock=0.0,
    train_phase=False,
    augmentations=None,
    sample_id=None,
):
    if train_phase:
        # We need the lambdas to make TensorFlow happy.
        # pylint: disable=unnecessary-lambda
        tf.cond(
            tf.math.not_equal(sample_rate, Config.audio_sample_rate),
            lambda: tf.Print(
                0,
                [
                    "WARNING: sample rate of sample",
                    sample_id,
                    "(",
                    sample_rate,
                    ") "
                    "does not match Config.audio_sample_rate. This can lead to incorrect results.",
                ],
            ),
            lambda: tf.no_op(),
            name="matching_sample_rate",
        )

    if train_phase and augmentations:
        audio = apply_graph_augmentations(
            "signal", audio, augmentations, transcript=transcript, clock=clock
        )

    spectrogram = contrib_audio.audio_spectrogram(
        audio,
        window_size=Config.audio_window_samples,
        stride=Config.audio_step_samples,
        magnitude_squared=True,
    )

    if train_phase and augmentations:
        spectrogram = apply_graph_augmentations(
            "spectrogram",
            spectrogram,
            augmentations,
            transcript=transcript,
            clock=clock,
        )

    features = contrib_audio.mfcc(
        spectrogram=spectrogram,
        sample_rate=sample_rate,
        dct_coefficient_count=Config.n_input,
        upper_frequency_limit=Config.audio_sample_rate / 2,
    )
    features = tf.reshape(features, [-1, Config.n_input])

    if train_phase and augmentations:
        features = apply_graph_augmentations(
            "features", features, augmentations, transcript=transcript, clock=clock
        )

    return features, tf.shape(input=features)[0]
//  [python]
def audio_to_features(
    audio,
    sample_rate,
    transcript=None,
    clock=0.0,
    train_phase=False,
    augmentations=None,
    sample_id=None,
):
    if train_phase:
        # We need the lambdas to make TensorFlow happy.
        # pylint: disable=unnecessary-lambda
        tf.cond(
            tf.math.not_equal(sample_rate, Config.audio_sample_rate),
            lambda: tf.print(
                "WARNING: sample rate of sample",
                sample_id,
                "(",
                sample_rate,
                ") "
                "does not match Config.audio_sample_rate. This can lead to incorrect results.",
            ),
            lambda: tf.no_op(),
            name="matching_sample_rate",
        )

    if train_phase and augmentations:
        audio = apply_graph_augmentations(
            "signal", audio, augmentations, transcript=transcript, clock=clock
        )

    spectrogram = contrib_audio.audio_spectrogram(
        audio,
        window_size=Config.audio_window_samples,
        stride=Config.audio_step_samples,
        magnitude_squared=True,
    )

    if train_phase and augmentations:
        spectrogram = apply_graph_augmentations(
            "spectrogram",
            spectrogram,
            augmentations,
            transcript=transcript,
            clock=clock,
        )

    features = contrib_audio.mfcc(
        spectrogram=spectrogram,
        sample_rate=sample_rate,
        dct_coefficient_count=Config.n_input,
        upper_frequency_limit=Config.audio_sample_rate / 2,
    )
    features = tf.reshape(features, [-1, Config.n_input])

    if train_phase and augmentations:
        features = apply_graph_augmentations(
            "features", features, augmentations, transcript=transcript, clock=clock
        )

    return features, tf.shape(input=features)[0]
//  [python]
def no_vars_loop(strat, iterable):
  for pr in iterable:
    t.print(strat.reduce('SUM', pr, axis=0))
//  [python]
def no_vars_loop(strat, iterable):
  for pr in iterable:
    tf.print(strat.reduce('SUM', pr, axis=0))
//  [python]
def _TrainModel(self, config):
    logging.info("Testing config:")
    logging.info(config)
    config = self._SetDefaults(config)
    self._ResetAllBackends()

    training_inputs, training_labels = (
        self._GetTrainingInputsAndLabels(config))

    keras_layer = cdf_layer.CDF(
        num_keypoints=config["num_keypoints"],
        units=config["units"],
        activation=config["activation"],
        reduction=config["reduction"],
        input_scaling_init=config["input_scaling_init"],
        input_scaling_type=config["input_scaling_type"],
        sparsity_factor=config["sparsity_factor"],
        kernel_initializer=config["kernel_initializer"],
        input_shape=(config["input_dims"],),
        dtype=tf.float32)
    model = keras.models.Sequential()
    model.add(keras_layer)

    # When we have multi-unit output, we average across the output units for
    # testing.
    if config["units"] > 1:
      model.add(
          keras.layers.Lambda(
              lambda x: tf.reduce_mean(x, axis=-1, keepdims=True)))

    optimizer = config["optimizer"](learning_rate=config["learning_rate"])
    model.compile(loss="mse", optimizer=optimizer)

    training_data = (training_inputs, training_labels)
    loss = test_utils.run_training_loop(
        config=config, training_data=training_data, keras_model=model
    )

    if tf.executing_eagerly():
      tf.compat.v1.logging.info("final weights: ", keras_layer.kernel)

    return loss
//  [python]
def _TrainModel(self, config):
    logging.info("Testing config:")
    logging.info(config)
    config = self._SetDefaults(config)
    self._ResetAllBackends()

    training_inputs, training_labels = (
        self._GetTrainingInputsAndLabels(config))

    keras_layer = cdf_layer.CDF(
        num_keypoints=config["num_keypoints"],
        units=config["units"],
        activation=config["activation"],
        reduction=config["reduction"],
        input_scaling_init=config["input_scaling_init"],
        input_scaling_type=config["input_scaling_type"],
        sparsity_factor=config["sparsity_factor"],
        kernel_initializer=config["kernel_initializer"],
        input_shape=(config["input_dims"],),
        dtype=tf.float32)
    model = keras.models.Sequential()
    model.add(keras_layer)

    # When we have multi-unit output, we average across the output units for
    # testing.
    if config["units"] > 1:
      model.add(
          keras.layers.Lambda(
              lambda x: tf.reduce_mean(x, axis=-1, keepdims=True)))

    optimizer = config["optimizer"](learning_rate=config["learning_rate"])
    model.compile(loss="mse", optimizer=optimizer)

    training_data = (training_inputs, training_labels)
    loss = test_utils.run_training_loop(
        config=config, training_data=training_data, keras_model=model
    )

    if tf.executing_eagerly():
      tf.print("final weights: ", keras_layer.kernel)

    return loss
//  [python]
def dataset_no_vars_loop(ds):
  for e in ds:
    t.print(e)
//  [python]
def dataset_no_vars_loop(ds):
  for e in ds:
    tf.print(e)
//  [python]
def iterator_no_vars_loop(ds):
  for e in iter(ds):
    t.print(e)
//  [python]
def iterator_no_vars_loop(ds):
  for e in iter(ds):
    tf.print(e)
//  [python]
def train_step(self, data):
    """The logic for one training step.

    This method is overridden to support custom training logic.
    This method is called by `Model.make_train_function`.
    This method should contain the mathematical logic for one step of training.
    This typically includes the forward pass, loss calculation, backpropagation,
    and metric updates.
    Configuration details for *how* this logic is run (e.g. `tf.function` and
    `tf.distribute.Strategy` settings), should be left to
    `Model.make_train_function`, which can also be overridden.

    Arguments:
      data: A dictionary of `Tensor`s.

    Returns:
      A `dict` containing values that will be passed to
      `tf.keras.callbacks.CallbackList.on_train_batch_end`. Typically, the
      values of the `Model`'s metrics are returned. Example:
      `{'loss': 0.2, 'accuracy': 0.7}`.
    """

    # If `sample_weight` is not provided, all samples will be weighted
    # equally.
    x, _, _ = tf.keras.utils.unpack_x_y_sample_weight(data)

    with tf.GradientTape() as tape:
      self(x, training=True)
      self.write_summary_on_train_batch_end()
      loss = self.loss_names_to_losses['total_loss']
      if tf.executing_eagerly():
        tf.print('total loss:', loss)
      trainable_variables = self.trainable_variables
      gradients = tape.gradient(loss, trainable_variables)
      self.optimizer.apply_gradients(zip(gradients, trainable_variables))

    # TODO(alirezafathi): rewrite following line to add metrics updates
    # self.compiled_metrics.update_state(y, y_pred, sample_weight)
    return {m.name: m.result() for m in self.metrics}
//  [python]
def _log_msg(cls, msg: str) -> None:
    if cls.output_method.lower().startswith("log"):
      logging.debug(msg)
    elif cls.output_method.lower() in ["tf", "tf.print", "tf_print"]:
      # So we don't have to import tensorflow if we aren't using it.
      import tensorflow as tf  # pylint: disable=g-import-not-at-top
      tf.print(msg)
    else:
      print(msg)
//  [python]
def fake_output_writer_op(input_list):
    return tf.print(input_list), tf.no_op