---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[34], line 211
    207 with strategy.scope():
    208     # Define and train the MAML model
    209     maml = MAML(maml_model, learning_rate=1e-4, meta_lr=1e-3, num_adapt_steps=5)
--> 211     maml.train(task_batches, epochs=100)
    213 # Predict on query set
    214 # prediction = predict_on_query_set(maml, X_train_norm, {'age': age_train, 'gender': gender_train_bin}, X_test_norm)

Cell In[34], line 123, in MAML.train(self, task_batches, epochs, validation_data, patience)
    121 for epoch in range(epochs):
    122     for task_data in task_batches:
--> 123         self.meta_train_step([task_data])  # Pass individual task data
    124     print(f"Epoch {epoch + 1}/{epochs} complete.")
    126     # Evaluate on validation data

Cell In[34], line 63, in MAML.meta_train_step(self, task_batch)
     61 # Apply gradients if they are not None
     62 if None not in meta_gradients:
---> 63     self.optimizer.apply_gradients(zip(meta_gradients, self.model.trainable_variables))

File /usr/local/lib/python3.10/site-packages/keras/src/optimizers/base_optimizer.py:291, in BaseOptimizer.apply_gradients(self, grads_and_vars)
    289 def apply_gradients(self, grads_and_vars):
    290     grads, trainable_variables = zip(*grads_and_vars)
--> 291     self.apply(grads, trainable_variables)
    292     # Return iterations for compat with tf.keras.
    293     return self._iterations

File /usr/local/lib/python3.10/site-packages/keras/src/optimizers/base_optimizer.py:356, in BaseOptimizer.apply(self, grads, trainable_variables)
    353     grads = [g if g is None else g / scale for g in grads]
    355 # Apply gradient updates.
--> 356 self._backend_apply_gradients(grads, trainable_variables)
    357 # Apply variable constraints after applying gradients.
    358 for variable in trainable_variables:

File /usr/local/lib/python3.10/site-packages/keras/src/optimizers/base_optimizer.py:419, in BaseOptimizer._backend_apply_gradients(self, grads, trainable_variables)
    416     self._apply_weight_decay(trainable_variables)
    418     # Run udpate step.
--> 419     self._backend_update_step(
    420         grads, trainable_variables, self.learning_rate
    421     )
    423 if self.use_ema:
    424     self._update_model_variables_moving_average(
    425         self._trainable_variables
    426     )

File /usr/local/lib/python3.10/site-packages/keras/src/backend/tensorflow/optimizer.py:121, in TFOptimizer._backend_update_step(self, grads, trainable_variables, learning_rate)
    119 grads_and_vars = list(zip(grads, trainable_variables))
    120 grads_and_vars = self._all_reduce_sum_gradients(grads_and_vars)
--> 121 tf.__internal__.distribute.interim.maybe_merge_call(
    122     self._distributed_tf_update_step,
    123     self._distribution_strategy,
    124     grads_and_vars,
    125     learning_rate,
    126 )

File /usr/local/lib/python3.10/site-packages/tensorflow/python/distribute/merge_call_interim.py:53, in maybe_merge_call(fn, strategy, *args, **kwargs)
     51   return fn(strategy, *args, **kwargs)
     52 else:
---> 53   return distribute_lib.get_replica_context().merge_call(
     54       fn, args=args, kwargs=kwargs)

AttributeError: 'NoneType' object has no attribute 'merge_call' class MAML:
    def __init__(self, model, learning_rate=1e-4, meta_lr=1e-3, num_adapt_steps=5):
        self.model = model
        self.learning_rate = learning_rate
        self.meta_lr = meta_lr
        self.num_adapt_steps = num_adapt_steps
        self.optimizer = tf.keras.optimizers.Adam(learning_rate=self.meta_lr)

#     def loss_function(self, y_true, y_pred):
#         """Compute combined loss for age and gender predictions."""
#         age_loss = tf.keras.losses.categorical_crossentropy(y_true['age'], y_pred[0])
#         gender_loss = tf.keras.losses.binary_crossentropy(y_true['gender'], y_pred[1])
#         return age_loss + gender_loss  # Combine losses
    def loss_function(self, y_true, y_pred):
        """Compute combined loss for age and gender predictions."""
        age_loss = tf.keras.losses.categorical_crossentropy(y_true['age'], y_pred[0])

        # Reshape y_true['gender'] to match y_pred[1]
        y_true_gender_reshaped = tf.expand_dims(y_true['gender'], axis=-1)  # Shape (8,) -> (8, 1)

        gender_loss = tf.keras.losses.binary_crossentropy(y_true_gender_reshaped, y_pred[1])
        return age_loss + gender_loss

    def adapt(self, x_support, y_support):
        for step in range(self.num_adapt_steps):
            initial_weights = [tf.identity(var) for var in self.model.trainable_variables]

            with tf.GradientTape() as tape:
                predictions = self.model(x_support, training=True)
                loss = self.loss_function(y_support, predictions)

            gradients = tape.gradient(loss, self.model.trainable_variables)

            # Log gradient norms
            gradient_norms = [tf.norm(grad).numpy() for grad in gradients if grad is not None]
            print(f"Adaptation Step {step + 1}, Gradient Norms: {gradient_norms}")

            # Ensure gradients are not None
            if None not in gradients:
                adapted_weights = [var - self.learning_rate * grad for var, grad in zip(self.model.trainable_variables, gradients)]

                for var, new_value in zip(self.model.trainable_variables, adapted_weights):
                    var.assign(new_value)

    def meta_train_step(self, task_batch):
        meta_gradients = [tf.zeros_like(weight) for weight in self.model.trainable_variables]

        for task_data in task_batch:
            x_support, y_support, x_query, y_query = task_data
            self.adapt(x_support, y_support)

            with tf.GradientTape() as tape:
                predictions = self.model(x_query, training=True)
                loss = self.loss_function(y_query, predictions)

            task_gradients = tape.gradient(loss, self.model.trainable_variables)

            if None not in task_gradients:
                meta_gradients = [meta_g + task_g for meta_g, task_g in zip(meta_gradients, task_gradients)]

        # Apply gradients if they are not None
        if None not in meta_gradients:
            self.optimizer.apply_gradients(zip(meta_gradients, self.model.trainable_variables))

#     def adapt(self, x_support, y_support):
#         """Perform inner-loop adaptation using the support set."""
#         for step in range(self.num_adapt_steps):
#             # Store the current weights before adaptation
#             initial_weights = [tf.identity(var) for var in self.model.trainable_variables]

#             with tf.GradientTape() as tape:
#                 predictions = self.model(x_support, training=True)
#                 loss = self.loss_function(y_support, predictions)

#             # Compute gradients for the current task (support set)
#             gradients = tape.gradient(loss, self.model.trainable_variables)

#             # Log the gradient norms for each layer
#             gradient_norms = [tf.norm(grad).numpy() for grad in gradients if grad is not None]
# #             print(f"Adaptation Step {step + 1}, Gradient Norms: {gradient_norms}")

#             # Update the model parameters using the task-specific learning rate
#             adapted_weights = [var - self.learning_rate * grad for var, grad in zip(self.model.trainable_variables, gradients)]

#             # Assign the adapted weights to the model variables
#             for var, new_value in zip(self.model.trainable_variables, adapted_weights):
#                 var.assign(new_value)

#             # Log the difference in weights after adaptation
#             weight_diffs = [tf.norm(var - init_var).numpy() for var, init_var in zip(self.model.trainable_variables, initial_weights)]
# #             print(f"Adaptation Step {step + 1}, Weight Differences: {weight_diffs}")

#     def meta_train_step(self, task_batch):
#         """Perform one meta-training step over a batch of tasks."""
#         meta_gradients = [tf.zeros_like(weight) for weight in self.model.trainable_variables]

#         for task_data in task_batch:
#             x_support, y_support, x_query, y_query = task_data  # Ensure exactly 4 elements

#             # Inner loop: Adapt the model to the current task (support set)
#             self.adapt(x_support, y_support)

#             # Evaluate on query set after adaptation
#             with tf.GradientTape() as tape:
#                 predictions = self.model(x_query, training=True)
#                 loss = self.loss_function(y_query, predictions)

#             # Compute the outer-loop gradients (meta-gradients)
#             task_gradients = tape.gradient(loss, self.model.trainable_variables)
#             meta_gradients = [meta_g + task_g for meta_g, task_g in zip(meta_gradients, task_gradients)]

#         # Apply the meta-gradient update to the model
#         self.optimizer.apply_gradients(zip(meta_gradients, self.model.trainable_variables))
        


    def train(self, task_batches, epochs, validation_data=None, patience=5):
        best_val_loss = float('inf')
        wait = 0
        """Train the model using the meta-learning approach."""
        for epoch in range(epochs):
            for task_data in task_batches:
                self.meta_train_step([task_data])  # Pass individual task data
            print(f"Epoch {epoch + 1}/{epochs} complete.")

            # Evaluate on validation data
            if validation_data:
                val_support_data, val_support_labels, val_query_data, val_query_labels = validation_data
                _, _, val_predictions = self.predict_query(val_support_data, val_support_labels, val_query_data)
                val_loss = self.loss_function(val_query_labels, val_predictions)
                print(f"Validation Loss: {val_loss.numpy()}")

                # Early stopping
                if val_loss < best_val_loss:
                    best_val_loss = val_loss
                    wait = 0
                else:
                    wait += 1
                    if wait >= patience:
                        print("Early stopping triggered.")
                        break


    def predict_query(self, x_support, y_support, x_query):
        """Adapt to the support set and predict the labels for the query set."""
        self.adapt(x_support, y_support)
        predictions = self.model.predict(x_query, verbose=2)
        age_predicted = np.argmax(predictions[0], axis=1)  # Age predictions
        gender_predicted = (predictions[1] > 0.5).astype(int)  # Gender predictions
        return age_predicted, gender_predicted, predictions


# Example function to adapt to support set and predict query set
def predict_on_query_set(maml, support_data, support_labels, query_data):
    """Adapt model on support set and predict query data."""
    # Adapt to support data
    print("Adapting model to support set...")
    _, _, predictions = maml.predict_query(support_data, support_labels, query_data)
    return predictions

def create_task_batch(train_data, train_labels, query_data, query_labels, batch_size=32):
    task_batch = []
    num_tasks = len(train_data) // batch_size

    for i in range(num_tasks):
        # Extract support set
        start_idx = i * batch_size
        end_idx = (i + 1) * batch_size
        x_support = train_data[start_idx:end_idx]
        
        # Slice the labels individually for age and gender
        y_support = {
            'age': train_labels['age'][start_idx:end_idx],
            'gender': train_labels['gender'][start_idx:end_idx]
        }

        # Ensure query indices do not exceed available data
        query_start_idx = (i + 1) * batch_size
        query_end_idx = min((i + 2) * batch_size, len(query_data))

        if query_start_idx >= len(query_data):
            break  # Stop if we exceed query data length

        # Extract query set
        x_query = query_data[query_start_idx:query_end_idx]
        y_query = {
            'age': query_labels['age'][query_start_idx:query_end_idx],
            'gender': query_labels['gender'][query_start_idx:query_end_idx]
        }

        # Append a tuple containing exactly 4 elements
        task_batch.append((x_support, y_support, x_query, y_query))

    return task_batch



# Example usage
X_query_norm = X_train_norm  # Your query data
age_query = age_train  # Your query age labels
gender_query_bin = gender_train_bin  # Your query gender labels

# Create task batches
task_batches = create_task_batch(X_train_norm, {'age': age_train, 'gender': gender_train_bin},
                                  X_query_norm, {'age': age_query, 'gender': gender_query_bin},
                                  batch_size=8)
with strategy.scope():
    # Define and train the MAML model
    maml = MAML(maml_model, learning_rate=1e-4, meta_lr=1e-3, num_adapt_steps=5)

    maml.train(task_batches, epochs=100)

# Predict on query set
# prediction = predict_on_query_set(maml, X_train_norm, {'age': age_train, 'gender': gender_train_bin}, X_test_norm)