Update m1_tf_test.py

m1_tf_test.py

@@ -8,10 +8,32 @@
 #############################################################
 import sys
 import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras.layers import Layer, Input
 tf.__version__
 tf.config.list_physical_devices()
 from random import randrange
 
+class ComputeSum(Layer):
+
+  def __init__(self, input_dim):
+      super(ComputeSum, self).__init__()
+      # Create a non-trainable weight.
+      self.total = tf.Variable(initial_value=tf.zeros((input_dim,)),
+                               trainable=True)
+
+  def call(self, inputs):
+      self.total.assign_add(tf.reduce_sum(inputs, axis=0))
+      return self.total
+
+def ComputeSumModel(input_shape):
+    inputs = Input(shape = input_shape)
+    outputs = ComputeSum(input_shape[0])(inputs)
+    
+    model = tf.keras.Model(inputs = inputs, outputs = outputs)
+
+    return model
+
 def xatoi(Str):
   
     sign, base, i = 1, 0, 0
@@ -47,7 +69,7 @@ inputs = tf.keras.Input(shape=(xatoi(sys.argv[1]),), name="digits")
 model = tf.keras.models.load_model('model')
 mnist = tf.keras.datasets.mnist
 (x_train, y_train), (x_test, y_test) = mnist.load_data()
-x_train, x_test = x_train / randrange(255), x_test / randrange(255)
+x_train, x_test = x_train / randrange(255+1), x_test / randrange(255+1)
 model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128,activation='selu',name='layer1'),
@@ -65,8 +87,11 @@ loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
 model.compile(optimizer='adam',
  loss=loss_fn,
  metrics=['accuracy'])
-model.fit(x_train, y_train, epochs=10)
-outputs = tf.keras.layers.Dense(4, activation='softmax', name='predictions')((x_train,y_train))
+these = model.fit(x_train, y_train, epochs=10)
+that = ComputeSum(len(these))
+those = that(these)
+outputs = tf.keras.layers.Dense(4, activation='softmax', name='predictions')(those)
+model = ComputeSumModel((tf.shape(these)[1],))
 model = tf.keras.Model(inputs=inputs, outputs=outputs)
 model.build()
 model.save('model')