Log gradients of keras model to tensorboard with distributed training

I'm training a tfrs model (thus a Keras model) using distributed training strategy supported by TF. Now at the end of each epoch, I want to log the gradients of some layers in the formats of (1) scalar (mean of gradients) and (2) histogram (gradients of all perceptions in that layer). For format (1) we can just return it as a metric so it's straightforward. However, for format (2) I don't find any way to make it work. I have tried the following approaches.

1. Use a callback function. But distributed strategy prevents gradients from being used outside of update functions, e.g. train_step. Also when gradients are computed in the train_step function, I cannot store it as a class variable cause outputs have to be in return values, see 1.

2. Directly log gradients in train_step. This works but I) it logs the gradients at each step instead of at the end of each epoch; II) I cannot get the correct step count inside the train_step function, i.e. the step count for all logged gradients are the same. Sample code is given at the following.

class MyModel(tfrs.models.Model):

... Some functions ...

def train_step(self, inputs):
  with tf.GradientTape(persistent=True, watch_accessed_variables=True) as tape:
          loss = self.compute_loss(inputs, training=True)
          regularization_loss = sum(self.losses)

          total_loss = loss + regularization_loss

        gradients = tape.gradient(total_loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

        grads = {}
        grads['A_layer'] = tape.gradient(loss, self.A_model.A_layers.trainable_weights)

        with self.w.as_default():
          for name in grads:
            g = grads[name]
            curr_grad = g[0]
            mean = tf.reduce_mean(tf.abs(curr_grad))
            tf.summary.scalar(f"grad_mean_layer_{name}", mean, step=self.epochs)
            tf.summary.histogram(f"grad_hist_layer_{name}", curr_grad, step=self.epochs)
          self.w.flush()
          self.epochs += 1