text: import argparse
  import tensorflow as tf
  from tensorflow.keras import layers
  import logictensornetworks as ltn
  import baselines, data, commons
  import argparse
  
  def parse_args():
      parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
      parser.add_argument('--csv-path',type=str,default="SDadd_ltn.csv")
      parser.add_argument('--epochs',type=int,default=20)
      parser.add_argument('--n-examples-train',type=int,default=30000)
      parser.add_argument('--n-examples-test',type=int,default=5000)
      parser.add_argument('--batch-size',type=int,default=32)
      args = parser.parse_args()
      dict_args = vars(args)
      return dict_args
  
  args = parse_args()
  n_examples_train = args['n_examples_train']
  n_examples_test = args['n_examples_test']
  batch_size = args['batch_size']
  EPOCHS = args['epochs']
  csv_path = args['csv_path']
  
  """ DATASET """
  
  ds_train, ds_test = data.get_mnist_op_dataset(
          count_train=n_examples_train,
          count_test=n_examples_test,
          buffer_size=10000,
          batch_size=batch_size,
          n_operands=2,
          op=lambda args: args[0]+args[1])
  
  """ LTN MODEL AND LOSS """
  ### Predicates
  logits_model = baselines.SingleDigit()
  Digit = ltn.Predicate(ltn.utils.LogitsToPredicateModel(logits_model))
  ### Variables
  d1 = ltn.variable("digits1", range(10))
  d2 = ltn.variable("digits2", range(10))
  ### Operators
  Not = ltn.Wrapper_Connective(ltn.fuzzy_ops.Not_Std())
  And = ltn.Wrapper_Connective(ltn.fuzzy_ops.And_Prod())
  Or = ltn.Wrapper_Connective(ltn.fuzzy_ops.Or_ProbSum())
  Implies = ltn.Wrapper_Connective(ltn.fuzzy_ops.Implies_Reichenbach())
  Forall = ltn.Wrapper_Quantifier(ltn.fuzzy_ops.Aggreg_pMeanError(),semantics="forall")
  Exists = ltn.Wrapper_Quantifier(ltn.fuzzy_ops.Aggreg_pMean(),semantics="exists")
  ### Axioms
  @tf.function
  def axioms(images_x, images_y, labels_z, p_schedule=tf.constant(2.)):
      images_x = ltn.variable("x", images_x)
      images_y = ltn.variable("y", images_y)
      labels_z = ltn.variable("z", labels_z)
      return Forall(
              ltn.diag(images_x,images_y,labels_z),
              Exists(
                  (d1,d2),
                  And(Digit([images_x,d1]),Digit([images_y,d2])),
                  mask_vars=[d1,d2,labels_z],
                  mask_fn=lambda vars: tf.equal(vars[0]+vars[1],vars[2]),
                  p=p_schedule
              ),
              p=2
          )
  
  """ TRAINING """
  
  optimizer = tf.keras.optimizers.Adam(0.001)
  metrics_dict = {
      'train_loss': tf.keras.metrics.Mean(name="train_loss"),
      'train_accuracy': tf.keras.metrics.Mean(name="train_accuracy"),
      'test_loss': tf.keras.metrics.Mean(name="test_loss"),
      'test_accuracy': tf.keras.metrics.Mean(name="test_accuracy")    
  }
  
  @tf.function
  def train_step(images_x, images_y, labels_z, **parameters):
      # loss
      with tf.GradientTape() as tape:
          loss = 1.- axioms(images_x, images_y, labels_z, **parameters)
      gradients = tape.gradient(loss, logits_model.trainable_variables)
      optimizer.apply_gradients(zip(gradients, logits_model.trainable_variables))
      metrics_dict['train_loss'](loss)
      # accuracy
      predictions_x = tf.argmax(logits_model(images_x),axis=-1)
      predictions_y = tf.argmax(logits_model(images_y),axis=-1)
      predictions_z = predictions_x + predictions_y
      match = tf.equal(predictions_z,tf.cast(labels_z,predictions_z.dtype))
      metrics_dict['train_accuracy'](tf.reduce_mean(tf.cast(match,tf.float32)))
      
  @tf.function
  def test_step(images_x, images_y, labels_z, **parameters):
      # loss
      loss = 1.- axioms(images_x, images_y, labels_z, **parameters)
      metrics_dict['test_loss'](loss)
      # accuracy
      predictions_x = tf.argmax(logits_model(images_x),axis=-1)
      predictions_y = tf.argmax(logits_model(images_y),axis=-1)
      predictions_z = predictions_x + predictions_y
      match = tf.equal(predictions_z,tf.cast(labels_z,predictions_z.dtype))
      metrics_dict['test_accuracy'](tf.reduce_mean(tf.cast(match,tf.float32)))
  
  from collections import defaultdict
  
  scheduled_parameters = defaultdict(lambda: {})
  for epoch in range(0,4):
      scheduled_parameters[epoch] = {"p_schedule":tf.constant(1.)}
  for epoch in range(4,8):
      scheduled_parameters[epoch] = {"p_schedule":tf.constant(2.)}
  for epoch in range(8,12):
      scheduled_parameters[epoch] = {"p_schedule":tf.constant(4.)}
  for epoch in range(12,EPOCHS):
      scheduled_parameters[epoch] = {"p_schedule":tf.constant(6.)}
  
  commons.train(
      EPOCHS,
      metrics_dict,
      ds_train,
      ds_test,
      train_step,
      test_step,
      csv_path=csv_path,
      scheduled_parameters=scheduled_parameters
  )