<jupyter_start><jupyter_text>Implementing Federated Averaging<jupyter_code>NUMBER_FEATURES = 784
NUMBER_TARGETS = 10<jupyter_output><empty_output><jupyter_text>Combining TensorFlow and TensorFlow Federated  Defining a loss function<jupyter_code>BATCH_SPEC = collections.OrderedDict(
    X=tf.TensorSpec(shape=[None, NUMBER_FEATURES], dtype=tf.float32),
    y=tf.TensorSpec(shape=[None], dtype=tf.int32)
)
BATCH_TYPE = tff.to_type(BATCH_SPEC)
print(BATCH_TYPE)
MODEL_SPEC = collections.OrderedDict(
    weights=tf.TensorSpec(shape=[NUMBER_FEATURES, NUMBER_TARGETS], dtype=tf.float32),
    bias=tf.TensorSpec(shape=[NUMBER_TARGETS], dtype=tf.float32)
)
MODEL_TYPE = tff.to_type(MODEL_SPEC)
print(MODEL_TYPE)
@tf.function
def linear_transformation(model, batch):
    W = model["weights"]
    b = model["bias"]
    X = batch['X']
    Z = tf.matmul(X, W) + b
    return Z


@tf.function
def crossentropy_loss(y_true, y_pred):
    crossentropy = tf.reduce_sum(tf.one_hot(y_true, NUMBER_TARGETS) * tf.math.log(y_pred),  axis=[1])
    return -tf.reduce_mean(crossentropy)


@tff.tf_computation(MODEL_TYPE, BATCH_TYPE)
@tf.function
def batch_loss(model, batch):
    y_true = batch['y']
    y_pred = tf.nn.softmax(linear_transformation(model, batch))
    return crossentropy_loss(y_true, y_pred)
print(batch_loss.type_signature)<jupyter_output>(<<weights=float32[784,10],bias=float32[10]>,<X=float32[?,784],y=int32[?]>> -> float32)<jupyter_text>Gradient descent on a single batch<jupyter_code>@tff.tf_computation(MODEL_TYPE, BATCH_TYPE, tf.float32)
def batch_train(initial_model, batch, learning_rate):
    # Define a group of model variables and set them to `initial_model`. Must
    # be defined outside the @tf.function.
    model_vars = collections.OrderedDict([
        (name, tf.Variable(name=name, initial_value=value))
        for name, value in initial_model.items()
    ])
    optimizer = tf.keras.optimizers.SGD(learning_rate)

    @tf.function
    def _train_on_batch(model_vars, batch):
        # Perform one step of gradient descent using loss from `batch_loss`.
        with tf.GradientTape() as tape:
            loss = forward_pass(model_vars, batch)
        grads = tape.gradient(loss, model_vars)
        optimizer.apply_gradients(zip(tf.nest.flatten(grads), tf.nest.flatten(model_vars)))
        return model_vars

    return _train_on_batch(model_vars, batch)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12, 'fraction_numerical': 0.015, 'mean_word_length': 5.252604166666667, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 16, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '174731', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 889, 'n_tokens_neox': 856, 'n_words': 147, 'path': 'notebooks/custom-federated-algorithms-part-2.ipynb', 'repo_name': 'kaust-vislab/tensorflow-federated-data-science-project', 'size': '2400'}