def test_accelerated_weighted_category_accuracy(self):
    """Test multi-device aggregation of weights."""
    layer = tl.Accelerate(tl.WeightedCategoryAccuracy())
    weights = np.array([1., 1., 1., 0.])
    targets = np.array([0, 1, 2, 3])

    model_outputs = np.array([[.2, .1, .7, 0.],
                              [.2, .1, .7, 0.],
                              [.2, .1, .7, 0.],
                              [.2, .1, .7, 0.]])
    accuracy = layer([model_outputs, targets, weights])
    self.assertEqual(np.mean(accuracy), 1 / 3)
Test multi-device aggregation of weights.
trax/layers/acceleration_test.py