Method fun

python/tests/ring_test_distributed.py:120–121

(x)

    def test_all_gather_vjp(self):
        def fun(x):
            return mx.distributed.all_gather(x)[0]

        dfdx = mx.grad(fun)(mx.array(1.0))
        if mx.distributed.init().rank() == 0: