jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def transfer(self, content_image, style_image, produce_image_file="produce_images", epochs=5, steps_per_epoch=100):
        if not os.path.exists(produce_image_file):
            os.mkdir(produce_image_file)
        self.content_image = content_image
        self.style_targets = self.extractor(style_image)['style']
        self.content_targets = self.extractor(content_image)['content']

        # Define a tf.Variable to contain the image to optimize.
        # To make this quick, initialize it with the content image
        # (the tf.Variable must be the same shape as the content image):
        produce_image = tf.Variable(self.content_image)

        start = time.time()
        step = 0
        for n in range(epochs):
            for m in range(steps_per_epoch):
                step += 1
                self.train_step(produce_image)
                tf.print(".", end='')
            self.VGG19_imshow(produce_image.read_value(),
                              title=f"{step}_steps", produce_image_file=produce_image_file)
        end = time.time()
        tf.print("Total time: {:.1f}".format(end - start))
//  [python]
def test(self, graph, mask, edge_set_name, masked_edge_set_name, expected):
    actual = ops.mask_edges(graph, edge_set_name, mask, masked_edge_set_name)

    for node_set_name in [*expected.node_sets, *actual.node_sets]:
      self.assertAllEqual(actual.node_sets[node_set_name].sizes,
                          expected.node_sets[node_set_name].sizes)
      self.assertAllEqual(actual.node_sets[node_set_name].features,
                          expected.node_sets[node_set_name].features)
    for edge_set_name in [*expected.edge_sets, *actual.edge_sets]:
      tf.print(actual.edge_sets[edge_set_name])
      self.assertAllEqual(actual.edge_sets[edge_set_name].sizes,
                          expected.edge_sets[edge_set_name].sizes)
      self.assertAllEqual(actual.edge_sets[edge_set_name].features,
                          expected.edge_sets[edge_set_name].features)
      self.assertAllEqual(actual.edge_sets[edge_set_name].adjacency.source,
                          expected.edge_sets[edge_set_name].adjacency.source)
      self.assertAllEqual(actual.edge_sets[edge_set_name].adjacency.target,
                          expected.edge_sets[edge_set_name].adjacency.target)
//  [python]
print('Image data type:', img.dtype)
print(img[100:102, 100:102, :])


# **INFO-BOX: The rank of a grayscale image for input to a CNN**



img_raw = tf.io.read_file('example-image-gray.png')
img = tf.image.decode_image(img_raw)
tf.print('Rank:', tf.rank(img))
//  [python]
def train_one_step(X, y):    
    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("loss", loss)
    grads = tape.gradient(loss, model.variables)    
    optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
//  [python]
def train_one_step(X, y):    
    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)
        # 注意這裡使用了TensorFlow內建的tf.print()。@tf.function不支援Python內建的print方法
        tf.print("loss", loss)
    grads = tape.gradient(loss, model.variables)    
    optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
//  [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, FLAGS.audio_sample_rate),
                lambda: tf.Print('WARNING: sample rate of sample', [sample_id, '(', sample_rate, ') ',
                                 'does not match FLAGS.audio_sample_rate. This can lead to incorrect results.'],
                                 summarize=1000),
                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=FLAGS.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, FLAGS.audio_sample_rate),
                lambda: tf.print('WARNING: sample rate of sample', sample_id, '(', sample_rate, ') '
                                 'does not match FLAGS.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=FLAGS.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 test_toy_rf_classification_winner_takes_all(self, add_boolean_features,
                                                  has_catset):

    # Create toy model.
    model_path = os.path.join(
        tempfile.mkdtemp(dir=self.get_temp_dir()), "test_basic_rf_wta")
    test_utils.build_toy_random_forest(
        model_path,
        winner_take_all_inference=True,
        add_boolean_features=add_boolean_features,
        has_catset=has_catset)
    features = test_utils.build_toy_input_feature_values(
        features=None, has_catset=has_catset)

    # Prepare model.
    model = inference.Model(model_path)

    @tf.function
    def init_model():
      tf.print("Loading model")
      model.init_op()

    @tf.function
    def apply_model(features):
      tf.print("Running model")
      return model.apply(features)

    init_model()

    predictions = apply_model(features)
    print("Predictions: %s", predictions)

    logging.info("dense_predictions_values: %s", predictions.dense_predictions)
    logging.info("dense_col_representation_values: %s",
                 predictions.dense_col_representation)

    expected_proba, expected_classes = test_utils.expected_toy_predictions_rf_wta(
        add_boolean_features=add_boolean_features, has_catset=has_catset)
    self.assertAllEqual(predictions.dense_col_representation, expected_classes)
    self.assertAllClose(predictions.dense_predictions, expected_proba)
//  [python]
def prepare_df(model, data, outdir, target, save_raw=False):
    print("prepare_df")

    dfs = []
    for iev, d in enumerate(data):
        if iev%50==0:
            tf.print(".", end="")
        X, y, w = d
        pred = model(X, training=False).numpy()
        pred_id_onehot, pred_charge, pred_momentum = separate_prediction(pred)
        pred_id = assign_label(pred_id_onehot).flatten()
 
        if save_raw:
            np.savez_compressed("ev_{}.npz".format(iev), X=X.numpy(), y=y.numpy(), w=w.numpy(), y_pred=pred)

        pred_charge = pred_charge[:, :, 0].flatten()
        pred_momentum = pred_momentum.reshape((pred_momentum.shape[0]*pred_momentum.shape[1], pred_momentum.shape[2]))

        true_id, true_charge, true_momentum = separate_truth(y)
        true_id = true_id.numpy()[:, :, 0].flatten()
        true_charge = true_charge.numpy()[:, :, 0].flatten()
        true_momentum = true_momentum.numpy().reshape((true_momentum.shape[0]*true_momentum.shape[1], true_momentum.shape[2]))
       
        df = pandas.DataFrame()
        df["pred_pid"] = np.array([int(class_labels[p]) for p in pred_id])
        df["pred_eta"] = np.array(pred_momentum[:, 0], dtype=np.float64)
        df["pred_phi"] = np.array(pred_momentum[:, 1], dtype=np.float64)
        df["pred_e"] = np.array(pred_momentum[:, 2], dtype=np.float64)

        df["{}_pid".format(target)] = np.array([int(class_labels[p]) for p in true_id])
        df["{}_eta".format(target)] = np.array(true_momentum[:, 0], dtype=np.float64)
        df["{}_phi".format(target)] = np.array(true_momentum[:, 1], dtype=np.float64)
        df["{}_e".format(target)] = np.array(true_momentum[:, 2], dtype=np.float64)

        df["iev"] = iev
        dfs += [df]
    df = pandas.concat(dfs, ignore_index=True)
    fn = outdir + "/df.pkl.bz2"
    df.to_pickle(fn)
    print("prepare_df done", fn)
//  [python]
a = tf.constant(1, name='a')
b = tf.constant(2, name='b')
c = tf.constant(3, name='c')

z = 2 * (a - b) + c

tf.print("my result: ", z)
//  [python]
def evaluate_constrained_model(model,
                               dataset: tf.Tensor,
                               constraints: psl_model.PSLModel,
                               grad_steps: int = 10,
                               alpha: float = 0.1) -> List[tf.Tensor]:
  """Custom evaluation step."""
  logits = []
  for x_batch, y_batch in dataset:
    batch_logits = test_step(
        model,
        x_batch,
        y_batch,
        constraints=constraints,
        grad_steps=grad_steps,
        alpha=alpha)
    logits.append(batch_logits)

  for metric in model.metrics:
    tf.print(metric, metric.result())

  return logits
//  [python]
def test_toy_rf_classification_winner_takes_all_v2(self, add_boolean_features,
                                                     has_catset):

    # Create toy model.
    model_path = os.path.join(
        tempfile.mkdtemp(dir=self.get_temp_dir()), "test_basic_rf_wta")
    test_utils.build_toy_random_forest(
        model_path,
        winner_take_all_inference=True,
        add_boolean_features=add_boolean_features,
        has_catset=has_catset)
    features = test_utils.build_toy_input_feature_values(
        features=None, has_catset=has_catset)

    # Prepare model.
    t.print("Loading model")
    model = inference.ModelV2(model_path)

    @tf.function
    def apply_non_eager(x):
      return model.apply(x)

    predictions_non_eager = apply_non_eager(features)
    predictions_eager = model.apply(features)

    def check_predictions(predictions):
      print("Predictions: %s", predictions)

      logging.info("dense_predictions_values: %s",
                   predictions.dense_predictions)
      logging.info("dense_col_representation_values: %s",
                   predictions.dense_col_representation)

      (expected_proba,
       expected_classes) = test_utils.expected_toy_predictions_rf_wta(
           add_boolean_features=add_boolean_features, has_catset=has_catset)
      self.assertAllEqual(predictions.dense_col_representation,
                          expected_classes)
      self.assertAllClose(predictions.dense_predictions, expected_proba)

    check_predictions(predictions_non_eager)
    check_predictions(predictions_eager)
//  [python]
def test_toy_rf_classification_winner_takes_all_v2(self, add_boolean_features,
                                                     has_catset):

    # Create toy model.
    model_path = os.path.join(
        tempfile.mkdtemp(dir=self.get_temp_dir()), "test_basic_rf_wta")
    test_utils.build_toy_random_forest(
        model_path,
        winner_take_all_inference=True,
        add_boolean_features=add_boolean_features,
        has_catset=has_catset)
    features = test_utils.build_toy_input_feature_values(
        features=None, has_catset=has_catset)

    # Prepare model.
    tf.print("Loading model")
    model = inference.ModelV2(model_path)

    @tf.function
    def apply_non_eager(x):
      return model.apply(x)

    predictions_non_eager = apply_non_eager(features)
    predictions_eager = model.apply(features)

    def check_predictions(predictions):
      print("Predictions: %s", predictions)

      logging.info("dense_predictions_values: %s",
                   predictions.dense_predictions)
      logging.info("dense_col_representation_values: %s",
                   predictions.dense_col_representation)

      (expected_proba,
       expected_classes) = test_utils.expected_toy_predictions_rf_wta(
           add_boolean_features=add_boolean_features, has_catset=has_catset)
      self.assertAllEqual(predictions.dense_col_representation,
                          expected_classes)
      self.assertAllClose(predictions.dense_predictions, expected_proba)

    check_predictions(predictions_non_eager)
    check_predictions(predictions_eager)
//  [python]
def create_model_fn(self):

    def model_fn(features, mode, config):
      if export_context.EXPORT_MODE != None:
        self._called_in_exported_mode = True

      table = hash_table_ops.test_hash_table(2)
      mtable = multi_hash_table_ops.MultiHashTable.from_configs(
          configs={"test": test_utils.generate_test_hash_table_config(1)})
      global_step = tf.compat.v1.train.get_or_create_global_step()
      if mode == tf.estimator.ModeKeys.PREDICT:
        output_tensor = table.lookup([0])
        output = tf.estimator.export.PredictOutput(output_tensor)
        moutput = tf.estimator.export.PredictOutput(
            mtable.lookup({"test": [0]})["test"])
        return tf.estimator.EstimatorSpec(
            mode=mode,
            predictions=output_tensor,
            export_outputs={
                tf.compat.v1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                    output,
                "table/lookup":
                    output,
                "mtable/lookup":
                    moutput,
            })
      add_op = table.assign_add([0], [[1, 2]]).as_op()
      add_op2 = mtable.assign_add({"test": ([0], [[1]])}).as_op()
      global_step = tf.compat.v1.assign_add(global_step, 1)
      with tf.control_dependencies([global_step]):
        print_op = tf.compat.v1.Print("tensor value:", table.lookup([0]))
        print_op2 = tf.compat.v1.Print("mtable tensor value: ",
                             mtable.lookup({"test": [0]}))
      ckpt_prefix = config.model_dir + "/model.ckpt"
      return tf.estimator.EstimatorSpec(
          mode=mode,
          train_op=tf.group([global_step, add_op, add_op2, print_op,
                             print_op2]),
          training_hooks=[
              tf.estimator.CheckpointSaverHook(
                  config.model_dir,
                  save_steps=1000,
                  listeners=[
                      hash_table_ops.HashTableCheckpointSaverListener(
                          ckpt_prefix),
                      multi_hash_table_ops.
                      MultiHashTableCheckpointSaverListener(ckpt_prefix),
                  ])
          ],
          loss=tf.constant(0.0))

    return model_fn
//  [python]
def create_model_fn(self):

    def model_fn(features, mode, config):
      if export_context.EXPORT_MODE != None:
        self._called_in_exported_mode = True

      table = hash_table_ops.test_hash_table(2)
      mtable = multi_hash_table_ops.MultiHashTable.from_configs(
          configs={"test": test_utils.generate_test_hash_table_config(1)})
      global_step = tf.compat.v1.train.get_or_create_global_step()
      if mode == tf.estimator.ModeKeys.PREDICT:
        output_tensor = table.lookup([0])
        output = tf.estimator.export.PredictOutput(output_tensor)
        moutput = tf.estimator.export.PredictOutput(
            mtable.lookup({"test": [0]})["test"])
        return tf.estimator.EstimatorSpec(
            mode=mode,
            predictions=output_tensor,
            export_outputs={
                tf.compat.v1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                    output,
                "table/lookup":
                    output,
                "mtable/lookup":
                    moutput,
            })
      add_op = table.assign_add([0], [[1, 2]]).as_op()
      add_op2 = mtable.assign_add({"test": ([0], [[1]])}).as_op()
      global_step = tf.compat.v1.assign_add(global_step, 1)
      with tf.control_dependencies([global_step]):
        print_op = tf.print("tensor value:", table.lookup([0]))
        print_op2 = tf.print("mtable tensor value: ",
                             mtable.lookup({"test": [0]}))
      ckpt_prefix = config.model_dir + "/model.ckpt"
      return tf.estimator.EstimatorSpec(
          mode=mode,
          train_op=tf.group([global_step, add_op, add_op2, print_op,
                             print_op2]),
          training_hooks=[
              tf.estimator.CheckpointSaverHook(
                  config.model_dir,
                  save_steps=1000,
                  listeners=[
                      hash_table_ops.HashTableCheckpointSaverListener(
                          ckpt_prefix),
                      multi_hash_table_ops.
                      MultiHashTableCheckpointSaverListener(ckpt_prefix),
                  ])
          ],
          loss=tf.constant(0.0))

    return model_fn
//  [python]
def print_metrics(metrics, step, every, decimals=2, name='metrics'):
  factor = 10 ** decimals
  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/{} ='.format(name, '/'.join(metrics.keys()))
    print_metrics = tf.cond(
        tf.equal(step % every, 0),
        lambda: tf.Print(message, [step] + [
            tf.round(mean.clear() * factor) / factor for mean in means]),
        tf.no_op)
  return print_metrics
//  [python]
def print_metrics(metrics, step, every, decimals=2, name='metrics'):
  factor = 10 ** decimals
  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/{} ='.format(name, '/'.join(metrics.keys()))
    print_metrics = tf.cond(
        tf.equal(step % every, 0),
        lambda: tf.print(message, [step] + [
            tf.round(mean.clear() * factor) / factor for mean in means]),
        tf.no_op)
  return print_metrics
//  [python]
def validate_model(self):
        with tf.device(self.player.device_name):
            with tf.name_scope("data_loading"):
                data_generator = self.test_dataset.generator_builder()
                data_iter = data_generator()

            # this maybe a bug in tensorflow
            # when using `model.fit` or `model.evalute` within tf.device context
            # program will stop
            with tf.name_scope("validate_model"):
                y_preds = []
                y_trues = []
                for index, (input_x, input_y) in enumerate(data_iter):
                    y_pred = self.model(input_x)
                    y_preds.append(y_pred)
                    y_trues.append(input_y)
                    if index + 1 >= self.test_dataset.iterations:
                        break

                y_pred = tf.concat(y_preds, axis=0)
                y_true = tf.concat(y_trues, axis=0)
                tf.print(self.metrics(y_true, y_pred))
//  [python]
a = tf.constant(1, name='a')
b = tf.constant(2, name='b')
c = tf.constant(3, name='c')

z = 2*(a - b) + c
tf.print('Result: z =', z)
//  [python]
def call(self, inputs: t_input):
        # tf.print(f"inputs at ComplexDense are {inputs.dtype}")
        if inputs.dtype != self.my_dtype:
            tf.print(f"WARNING: {self.name} - Expected input to be {self.my_dtype}, but received {inputs.dtype}.")
            if self.my_dtype.is_complex and inputs.dtype.is_floating:
                tf.print("\tThis is normally fixed using ComplexInput() "
                         "at the start (tf casts input automatically to real).")
            inputs = tf.cast(inputs, self.my_dtype)
        if self.my_dtype.is_complex:
            w = tf.complex(self.w_r, self.w_i)
            if self.use_bias:
                b = tf.complex(self.b_r, self.b_i)
        else:
            w = self.w
            if self.use_bias:
                b = self.b
        out = tf.matmul(inputs, w)
        if self.use_bias:
            out = out + b
        return self.activation(out)
//  [python]
def on_train_begin(self, logs: dict):
        print("Training starting time noted.")
        self.train_start_time = timer()
//  [python]
def on_train_begin(self, logs: dict):
        tf.print("Training starting time noted.", output_stream=sys.stdout)
        self.train_start_time = timer()
//  [python]
def debug(x: PondTensor, summarize=None, message=""):
    """Print contents of a PondTensor for debugging purposes."""
    if isinstance(x, PondPublicTensor):
        t.print(
            x.value_on_0.value,
            [x.value_on_0.value],
            summarize=summarize,
            message=message,
        )

    elif isinstance(x, PondPrivateTensor):
        t.print(
            x.share0.value,
            [x.reveal().value_on_0.value],
            summarize=summarize,
            message=message,
        )

    else:
        raise TypeError("Don't know how to debug {}".format(type(x)))
//  [python]
def debug(x: PondTensor, summarize=None, message=""):
    """Print contents of a PondTensor for debugging purposes."""
    if isinstance(x, PondPublicTensor):
        tf.print(
            x.value_on_0.value,
            [x.value_on_0.value],
            summarize=summarize,
            message=message,
        )

    elif isinstance(x, PondPrivateTensor):
        tf.print(
            x.share0.value,
            [x.reveal().value_on_0.value],
            summarize=summarize,
            message=message,
        )

    else:
        raise TypeError("Don't know how to debug {}".format(type(x)))
//  [python]
def on_train_end(self, logs: dict):
        t.print("Training ending time noted.", output_stream=sys.stdout)
        self.train_end_time = timer()
//  [python]
def on_train_end(self, logs: dict):
        tf.print("Training ending time noted.", output_stream=sys.stdout)
        self.train_end_time = timer()
//  [python]
def calc_rb_mask_decorator(fn):
    def wrapper(*args, **kwargs):
        if tf.distribute.get_replica_context():
            thread_id = tf.distribute.get_replica_context().replica_id_in_sync_group
        else:
            thread_id = "dummy"

        tf.compat.v1.print(thread_id, args[1], output_stream=f"file://{MASKS_SEEDS_PATH}")
        mask = fn(*args, **kwargs)
        return mask

    return wrapper
//  [python]
def calc_rb_mask_decorator(fn):
    def wrapper(*args, **kwargs):
        if tf.distribute.get_replica_context():
            thread_id = tf.distribute.get_replica_context().replica_id_in_sync_group
        else:
            thread_id = "dummy"

        tf.print(thread_id, args[1], output_stream=f"file://{MASKS_SEEDS_PATH}")
        mask = fn(*args, **kwargs)
        return mask

    return wrapper
//  [python]
def return_with_default(x):
  if x > 0:
    t.print('x', x)
    return x
  return x * x
//  [python]
def return_with_default(x):
  if x > 0:
    tf.print('x', x)
    return x
  return x * x
//  [python]
def call(self, inputs, training=None):
        if inputs.dtype != self.my_dtype:
            tf.print(f"Warning: Expecting input dtype {self.my_dtype} but got {inputs.dtype}. "
                     f"Automatic cast will be done.")
            inputs = tf.cast(inputs, dtype=self.my_dtype)
        if training is None:
            training = K.learning_phase()
            tf.print(f"Training was None and now is {training}")
            # This is used for my own debugging, I don't know WHEN this happens,
            # I trust K.learning_phase() returns a correct boolean.
        if training:
            # First get the mean and var
            mean = tf.math.reduce_mean(inputs, axis=self.used_axis)
            if self.cov_method == 1:
                X_20 = tf.concat((tf.math.real(inputs), tf.math.imag(inputs)), axis=-1)
                var_20_20 = tfp.stats.covariance(X_20, sample_axis=self.used_axis, event_axis=-1)
                valu = int(var_20_20.shape[-1] / 2)
                indices = [([[i, i], [i, i + valu]], [[i + valu, i], [i + valu, i + valu]]) for i in range(0, valu)]
                var = tf.gather_nd(var_20_20, indices=indices)
            elif self.cov_method == 2:
                X_10_2 = tf.stack((tf.math.real(inputs), tf.math.imag(inputs)), axis=-1)
                var_10_2_2 = tfp.stats.covariance(X_10_2, sample_axis=self.used_axis, event_axis=-1)
                var = var_10_2_2
            else:
                raise ValueError(f"Method {self.method} not implemented")

            # Now the train part with these values
            self.moving_mean.assign(self.momentum * self.moving_mean + (1. - self.momentum) * mean)
            self.moving_var.assign(self.moving_var * self.momentum + var * (1. - self.momentum))
            out = self._normalize(inputs, var, mean)
        else:
            out = self._normalize(inputs, self.moving_var, self.moving_mean)
        if self.scale:
            if self.my_dtype.is_complex:
                gamma = tf.complex(self.gamma_r, self.gamma_i)      # TODO: Should this be real valued?
            else:
                gamma = self.gamma
            out = gamma * out
        if self.center:
            if self.my_dtype.is_complex:
                beta = tf.complex(self.beta_r, self.beta_i)
            else:
                beta = self.beta
            out = out + beta
        return out
//  [python]
def print_layer_weights_and_shape(layer):
  """Prints variables information corresponding to a Keras layer.

  This function prints the name and the shape of its associated weights
  of all variables (trainable and untrainable) in a Keras layer.

  Args:
    layer: A Keras.layer.Layer object
  """
  weights = layer.get_weights()
  variables = layer.variables

  for i in range(len(weights)):
    tf.print(np.shape(weights[i]), variables[i].name)
//  [python]
def train_model(perm, best_valid_loss, old_valid_loss, i_worse):
    for i in range(n_max_epochs):

        perm = tf.random.shuffle(perm)
        for j in range(0, len(D_train_tf), batch_size):
            lst = [q for q in range(j, j + batch_size)]
            indices = tf.gather(perm, lst).numpy()
            step(indices)

        # Validation stage
        val_loss = calc_val_loss(E_val_tf, F_val_tf, D_val_tf)
        if val_loss < best_valid_loss:
            model.save('model')
            best_valid_loss = val_loss
        if val_loss >= old_valid_loss:
            i_worse += 1
        else:
            i_worse = 0
        if i_worse > patience:
            tf.print("Early stopping at Epoch {}".format(i))
            break
        old_valid_loss = val_loss
        if i % 500 == 0:
            print("Finished epoch: {} with val_loss: {}".format(i, val_loss))
//  [python]
def call(self, inputs, training=True):
    voxel_features, voxel_xyz_indices, num_valid_voxels = inputs
    voxel_features_list = [voxel_features]
    voxel_xyz_indices_list = [voxel_xyz_indices]
    num_valid_voxels_list = [num_valid_voxels]
    index_mapping_list = []

    if training:
      tf.print('num_valid_voxels', num_valid_voxels)

    # Encoder
    for level in range(self.num_levels):
      inputs_i = (voxel_features_list[-1], voxel_xyz_indices_list[-1],
                  num_valid_voxels_list[-1])
      conv_block_i = getattr(self, 'encoder_' + str(level))
      outputs_i = conv_block_i(inputs_i, training)
      (pooled_voxel_features, pooled_voxel_indices, num_valid_pooled_voxels,
       index_mapping) = sparse_voxel_net_utils.voxel_pooling(
           voxel_features=outputs_i,
           voxel_xyz_indices=inputs_i[1],
           num_valid_voxels=inputs_i[2],
           pooling_size=(2, 2, 2),
           segment_func=self.network_pooling_segment_func)
      voxel_features_list.append(pooled_voxel_features)
      voxel_xyz_indices_list.append(pooled_voxel_indices)
      num_valid_voxels_list.append(num_valid_pooled_voxels)
      index_mapping_list.append(index_mapping)

    # Bottleneck
    outputs_midl = self.middle_layer((
        voxel_features_list[-1],
        voxel_xyz_indices_list[-1],
        num_valid_voxels_list[-1],
    ), training)
    voxel_features_list[-1] = outputs_midl

    # Decoder
    for level in reversed(range(self.num_levels)):
      unpooled_features = sparse_voxel_net_utils.voxel_upsampling(
          pooled_voxel_features=voxel_features_list[level + 1],
          index_mapping=index_mapping_list[level])
      concatenated_features = tf.concat(
          [voxel_features_list[level], unpooled_features], axis=2)
      inputs_i = (concatenated_features, voxel_xyz_indices_list[level],
                  num_valid_voxels_list[level])
      conv_block_i = getattr(self, 'decoder_' + str(level))
      outputs_i = conv_block_i(inputs_i, training)
      voxel_features_list[level] = outputs_i

    outputs = {}
    # Output head convolutions
    for task_name in sorted(self.task_names_to_num_output_channels):
      conv_block_task_1 = getattr(self, f'{task_name}/final_conv1_block')
      conv_block_task_2 = getattr(self, f'{task_name}/final_conv2_block')
      net = conv_block_task_1(
          (voxel_features_list[0], voxel_xyz_indices, num_valid_voxels_list[0]),
          training)
      outputs[task_name] = conv_block_task_2(
          (net, voxel_xyz_indices, num_valid_voxels_list[0]), training)

    return outputs
//  [python]
def evaluate(model, tf_dataset, max_step=1000, num_eval_samples=None):
  """Evaluate the model."""
  if num_eval_samples:
    num_samples = num_eval_samples
  elif FLAGS.num_eval_samples:
    num_samples = FLAGS.num_eval_samples
  else:
    num_samples = 1
  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 _add_group_logs(self, cur_group_loss, step_count):
    """Add to group loss and weights to tensorboard."""
    tf.print("step_count:", step_count, output_stream=sys.stdout)
    tf.print("group loss:", cur_group_loss, output_stream=sys.stdout)
    tf.print("group est. loss:", self._group_loss, output_stream=sys.stdout)
    tf.print("group weights:", self._group_weights, output_stream=sys.stdout)
    group_summary = {}
    group_loss_summary = {
        f"Batch_GroupLoss_{self._group_labels[i]}": cur_group_loss[i]
        for i in range(self._num_groups)
    }
    group_loss_est_summary = {
        f"Est_GroupLoss_{self._group_labels[i]}": self._group_loss[i]
        for i in range(self._num_groups)
    }
    group_weight_summary = {
        f"GroupWeights_{self._group_labels[i]}": self._group_weights[i]
        for i in range(self._num_groups)
    }
    group_summary.update(group_loss_summary)
    group_summary.update(group_loss_est_summary)
    group_summary.update(group_weight_summary)
    self._add_tf_scalar_summary(group_summary, step_count)
//  [python]
def testPrintFormat(self):
    observation = tf.constant(-1)
    reward = tf.constant(2.0)

    time_step = ts.transition(observation, reward)
    with self.captureWritesToStream(sys.stdout) as printed:
      print_op = tf.compat.v1.print(time_step, output_stream=sys.stdout)
      self.evaluate(print_op)

    expected = "{'step_type': 1, 'reward': 2, 'discount': 1, 'observation': -1}"
    self.assertIn(expected, printed.contents())
//  [python]
def testPrintFormat(self):
    observation = tf.constant(-1)
    reward = tf.constant(2.0)

    time_step = ts.transition(observation, reward)
    with self.captureWritesToStream(sys.stdout) as printed:
      print_op = tf.print(time_step, output_stream=sys.stdout)
      self.evaluate(print_op)

    expected = "{'step_type': 1, 'reward': 2, 'discount': 1, 'observation': -1}"
    self.assertIn(expected, printed.contents())
//  [python]
def call(self, inputs, **kwargs):
        # Check that input_ids starts with control token
        if isinstance(inputs, (tuple, list)):
            input_ids = inputs[0]
        elif isinstance(inputs, dict):
            input_ids = inputs.get('input_ids')
        else:
            input_ids = inputs

        if tf.not_equal(tf.reduce_sum(input_ids[:, 0]), 0):
            tf.compat.v1.print("A sequence with no special tokens has been passed to the RoBERTa model. "
                           "This model requires special tokens in order to work. "
                           "Please specify add_special_tokens=True in your encoding.")

        return super(TFRobertaMainLayer, self).call(inputs, **kwargs)
//  [python]
def call(self, inputs, **kwargs):
        # Check that input_ids starts with control token
        if isinstance(inputs, (tuple, list)):
            input_ids = inputs[0]
        elif isinstance(inputs, dict):
            input_ids = inputs.get('input_ids')
        else:
            input_ids = inputs

        if tf.not_equal(tf.reduce_sum(input_ids[:, 0]), 0):
            tf.print("A sequence with no special tokens has been passed to the RoBERTa model. "
                           "This model requires special tokens in order to work. "
                           "Please specify add_special_tokens=True in your encoding.")

        return super(TFRobertaMainLayer, self).call(inputs, **kwargs)
//  [python]
def create_model_fn(self):

    def model_fn(features, mode, config, **kwargs):
      slot = self.ctx.feature_factory.create_feature_slot(
          feature.FeatureSlotConfig(name="slot"))
      s = slot.add_feature_slice(5)
      fc = feature.FeatureColumnV1(slot, "feature")
      embedding = fc.embedding_lookup(s)
      if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode, predictions=tf.constant(0))
      all_embeddings = [fc.get_all_embeddings_concat()]
      loss = tf.reduce_sum(embedding)
      grads = tf.gradients(loss, all_embeddings)
      print1 = tf.Print("embedding: ", [embedding])
      print2 = tf.Print("all_embeddings: ", [all_embeddings])
      with tf.control_dependencies([print1, print2]):
        train_op = tf.group(
            self._ctx.feature_factory.apply_gradients(zip(
                grads, all_embeddings)))
      return tf.estimator.EstimatorSpec(mode,
                                        train_op=train_op,
                                        loss=loss,
                                        predictions=tf.constant(0))

    return model_fn
//  [python]
def create_model_fn(self):

    def model_fn(features, mode, config, **kwargs):
      slot = self.ctx.feature_factory.create_feature_slot(
          feature.FeatureSlotConfig(name="slot"))
      s = slot.add_feature_slice(5)
      fc = feature.FeatureColumnV1(slot, "feature")
      embedding = fc.embedding_lookup(s)
      if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode, predictions=tf.constant(0))
      all_embeddings = [fc.get_all_embeddings_concat()]
      loss = tf.reduce_sum(embedding)
      grads = tf.gradients(loss, all_embeddings)
      print1 = tf.print("embedding: ", embedding)
      print2 = tf.print("all_embeddings: ", all_embeddings)
      with tf.control_dependencies([print1, print2]):
        train_op = tf.group(
            self._ctx.feature_factory.apply_gradients(zip(
                grads, all_embeddings)))
      return tf.estimator.EstimatorSpec(mode,
                                        train_op=train_op,
                                        loss=loss,
                                        predictions=tf.constant(0))

    return model_fn
//  [python]
def drop_path(net, keep_prob, is_training=True):
  """Drops out a whole example hiddenstate with the specified probability."""
  if is_training:
    batch_size = tf.shape(net)[0]
    noise_shape = [batch_size, 1, 1, 1]
    keep_prob = tf.cast(keep_prob, dtype=net.dtype)
    random_tensor = keep_prob
    random_tensor += tf.random_uniform(noise_shape, dtype=net.dtype)
    binary_tensor = tf.floor(random_tensor)
    net = tf.div(net, keep_prob) * binary_tensor
  return net
//  [python]
def drop_path(net, keep_prob, is_training=True):
  """Drops out a whole example hiddenstate with the specified probability."""
  if is_training:
    batch_size = tf.shape(net)[0]
    noise_shape = [batch_size, 1, 1, 1]
    keep_prob = tf.cast(keep_prob, dtype=net.dtype)
    random_tensor = keep_prob
    random_tensor += tf.random_uniform(noise_shape, dtype=net.dtype)
    binary_tensor = tf.floor(random_tensor)
    net = tf.divide(net, keep_prob) * binary_tensor
  return net
//  [python]
def drop_connect(inputs, is_training, survival_prob):
  """Drop the entire conv with given survival probability."""
  # "Deep Networks with Stochastic Depth", https://arxiv.org/pdf/1603.09382.pdf
  if not is_training:
    return inputs

  # Compute tensor.
  batch_size = tf.shape(inputs)[0]
  random_tensor = survival_prob
  random_tensor += tf.random_uniform([batch_size, 1, 1, 1], dtype=inputs.dtype)
  binary_tensor = tf.floor(random_tensor)
  # Unlike conventional way that multiply survival_prob at test time, here we
  # divide survival_prob at training time, such that no addition compute is
  # needed at test time.
  output = tf.div(inputs, survival_prob) * binary_tensor
  return output
//  [python]
def drop_connect(inputs, is_training, survival_prob):
  """Drop the entire conv with given survival probability."""
  # "Deep Networks with Stochastic Depth", https://arxiv.org/pdf/1603.09382.pdf
  if not is_training:
    return inputs

  # Compute tensor.
  batch_size = tf.shape(inputs)[0]
  random_tensor = survival_prob
  random_tensor += tf.random_uniform([batch_size, 1, 1, 1], dtype=inputs.dtype)
  binary_tensor = tf.floor(random_tensor)
  # Unlike conventional way that multiply survival_prob at test time, here we
  # divide survival_prob at training time, such that no addition compute is
  # needed at test time.
  output = tf.divide(inputs, survival_prob) * binary_tensor
  return output
//  [python]
def softmax(self, target, axis, name=None):
        max_axis = tf.reduce_max(target, axis, keepdims=True)
        target_exp = tf.exp(target-max_axis)
        normalize = tf.reduce_sum(target_exp, axis, keepdims=True)
        softmax = tf.div(target_exp, normalize, name)
        return softmax
//  [python]
def softmax(self, target, axis, name=None):
        max_axis = tf.reduce_max(target, axis, keepdims=True)
        target_exp = tf.exp(target-max_axis)
        normalize = tf.reduce_sum(target_exp, axis, keepdims=True)
        softmax = tf.divide(target_exp, normalize, name)
        return softmax
//  [python]
def _network_template(self, state):
    """Builds the convolutional network used to compute the agent's Q-values.

    Args:
      state: `tf.Tensor`, contains the agent's current state.

    Returns:
      net: _network_type object containing the tensors output by the network.
    """
    weights_initializer = slim.variance_scaling_initializer(
        factor=1.0 / np.sqrt(3.0), mode='FAN_IN', uniform=True)

    net = tf.cast(state, tf.float32)
    net = tf.div(net, 255.)
    net = slim.conv2d(
        net, 32, [8, 8], stride=4, weights_initializer=weights_initializer)
    net = slim.conv2d(
        net, 64, [4, 4], stride=2, weights_initializer=weights_initializer)
    net = slim.conv2d(
        net, 64, [3, 3], stride=1, weights_initializer=weights_initializer)
    net = slim.flatten(net)
    net = noisy_dqn_agent.fully_connected(
        net, 512, scope='fully_connected')
    net = noisy_dqn_agent.fully_connected(
        net,
        self.num_actions * self._num_atoms,
        activation_fn=None,
        scope='fully_connected_1')

    logits = tf.reshape(net, [-1, self.num_actions, self._num_atoms])
    # probabilities = contrib_layers.softmax(logits)
    probabilities = slim.softmax(logits)
    q_values = tf.reduce_sum(self._support * probabilities, axis=2)
    return self._get_network_type()(q_values, logits, probabilities)
//  [python]
def _network_template(self, state):
    """Builds the convolutional network used to compute the agent's Q-values.

    Args:
      state: `tf.Tensor`, contains the agent's current state.

    Returns:
      net: _network_type object containing the tensors output by the network.
    """
    weights_initializer = slim.variance_scaling_initializer(
        factor=1.0 / np.sqrt(3.0), mode='FAN_IN', uniform=True)

    net = tf.cast(state, tf.float32)
    net = tf.divide(net, 255.)
    net = slim.conv2d(
        net, 32, [8, 8], stride=4, weights_initializer=weights_initializer)
    net = slim.conv2d(
        net, 64, [4, 4], stride=2, weights_initializer=weights_initializer)
    net = slim.conv2d(
        net, 64, [3, 3], stride=1, weights_initializer=weights_initializer)
    net = slim.flatten(net)
    net = noisy_dqn_agent.fully_connected(
        net, 512, scope='fully_connected')
    net = noisy_dqn_agent.fully_connected(
        net,
        self.num_actions * self._num_atoms,
        activation_fn=None,
        scope='fully_connected_1')

    logits = tf.reshape(net, [-1, self.num_actions, self._num_atoms])
    # probabilities = contrib_layers.softmax(logits)
    probabilities = slim.softmax(logits)
    q_values = tf.reduce_sum(self._support * probabilities, axis=2)
    return self._get_network_type()(q_values, logits, probabilities)
//  [python]
def create_model_fn(self):

    def model_fn(features, mode, config):
      slot = self.ctx.feature_factory.create_feature_slot(
          feature.FeatureSlotConfig(
              name="slot",
              default_vec_initializer=entry.ConstantsInitializer(0),
              default_vec_optimizer=entry.AdagradOptimizer(
                  learning_rate=0.1, initial_accumulator_value=1)))
      s = slot.add_feature_slice(5)
      fc = feature.FeatureColumnV1(slot, "feature")
      embedding = fc.embedding_lookup(s)
      tf_embeddings = tf.Variable(initial_value=tf.zeros(shape=(4, 5)),
                                  name='embedding')
      tf_embedding = tf.reduce_sum(tf.nn.embedding_lookup(
          params=tf_embeddings, ids=features["tf_feature"]),
                                   axis=1)
      if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode, predictions=tf.constant(0))
      all_embeddings = [fc.get_all_embeddings_concat()]
      loss = tf.reduce_sum(embedding)
      tf_loss = tf.reduce_sum(tf_embedding)
      grads = tf.gradients(loss, all_embeddings)
      tf_grads = tf.gradients(tf_loss, tf_embeddings)
      gs = tf.compat.v1.train.get_or_create_global_step()
      print1 = tf.compat.v1.print(gs, "embedding: ", embedding)
      print2 = tf.compat.v1.print(gs, "tf_embedding: ", tf_embedding)
      assert_equal = tf.compat.v1.assert_equal(embedding, tf_embedding)
      with tf.control_dependencies([print1, print2, assert_equal]):
        train_op = tf.group(
            self._ctx.feature_factory.apply_gradients(zip(
                grads, all_embeddings)),
            tf.compat.v1.train.AdagradOptimizer(
                learning_rate=0.1, initial_accumulator_value=1).apply_gradients(
                    zip(tf_grads, [tf_embeddings])), gs.assign_add(1))
      return tf.estimator.EstimatorSpec(mode,
                                        train_op=train_op,
                                        loss=loss + tf_loss,
                                        predictions=tf.constant(0))

    return model_fn
//  [python]
def create_model_fn(self):

    def model_fn(features, mode, config):
      slot = self.ctx.feature_factory.create_feature_slot(
          feature.FeatureSlotConfig(
              name="slot",
              default_vec_initializer=entry.ConstantsInitializer(0),
              default_vec_optimizer=entry.AdagradOptimizer(
                  learning_rate=0.1, initial_accumulator_value=1)))
      s = slot.add_feature_slice(5)
      fc = feature.FeatureColumnV1(slot, "feature")
      embedding = fc.embedding_lookup(s)
      tf_embeddings = tf.Variable(initial_value=tf.zeros(shape=(4, 5)),
                                  name='embedding')
      tf_embedding = tf.reduce_sum(tf.nn.embedding_lookup(
          params=tf_embeddings, ids=features["tf_feature"]),
                                   axis=1)
      if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode, predictions=tf.constant(0))
      all_embeddings = [fc.get_all_embeddings_concat()]
      loss = tf.reduce_sum(embedding)
      tf_loss = tf.reduce_sum(tf_embedding)
      grads = tf.gradients(loss, all_embeddings)
      tf_grads = tf.gradients(tf_loss, tf_embeddings)
      gs = tf.compat.v1.train.get_or_create_global_step()
      print1 = tf.print(gs, "embedding: ", embedding)
      print2 = tf.print(gs, "tf_embedding: ", tf_embedding)
      assert_equal = tf.compat.v1.assert_equal(embedding, tf_embedding)
      with tf.control_dependencies([print1, print2, assert_equal]):
        train_op = tf.group(
            self._ctx.feature_factory.apply_gradients(zip(
                grads, all_embeddings)),
            tf.compat.v1.train.AdagradOptimizer(
                learning_rate=0.1, initial_accumulator_value=1).apply_gradients(
                    zip(tf_grads, [tf_embeddings])), gs.assign_add(1))
      return tf.estimator.EstimatorSpec(mode,
                                        train_op=train_op,
                                        loss=loss + tf_loss,
                                        predictions=tf.constant(0))

    return model_fn
//  [python]
def tf_matmul_separate(A,B):
  batch,a_h,a_w,_,_ = A.get_shape()
  if batch*a_h*a_w < 80000:
      return tf.matmul(A,B)
  w_half = tf.cast(tf.div(a_w,2),tf.int32)
  h_half = tf.cast(tf.div(a_h,2),tf.int32)
  A_1 = A[:,:h_half,:w_half,:,:]
  A_2 = A[:,:h_half,w_half:,:,:]
  A_3 = A[:,h_half:,:w_half,:,:]
  A_4 = A[:,h_half:,w_half:,:,:]
  
  B_1 = B[:,:h_half,:w_half,:,:]
  B_2 = B[:,:h_half,w_half:,:,:]
  B_3 = B[:,h_half:,:w_half,:,:]
  B_4 = B[:,h_half:,w_half:,:,:]
  
  res_1 = tf_matmul_separate(A_1,B_1)
  res_2 = tf_matmul_separate(A_2,B_2)
  res_3 = tf_matmul_separate(A_3,B_3)
  res_4 = tf_matmul_separate(A_4,B_4)
  tmp_up = tf.concat([res_1,res_2],2)
  tmp_down = tf.concat([res_3,res_4],2)
  final_res = tf.concat([tmp_up,tmp_down],1)
  return final_res