import torch


def aggregate(gradients, f, p=0.9, **kwargs):
  """ NaN-resilient median coordinate-per-coordinate rule.
  Args:
    gradients Non-empty list of gradients to aggregate
    f         Number of Byzantine gradients to tolerate
    p         Median selection probability
    ...       Ignored keyword-arguments
  Returns:
    NaN-resilient, coordinate-wise median of the gradients
  """
  # Sample selection indications
  c = torch.distributions.bernoulli.Bernoulli(torch.ones_like(gradients[0]) * p).sample()
  # Compute median and mask according to selection
  m = torch.stack(gradients).median(dim=0)[0].mul_(c)
  # Inverse selection
  c.neg_().add_(1)
  # Add masked first gradient and return
  return m.add_(gradients[0].mul(c))