# Copyright © 2023 Apple Inc. import unittest import mlx.core as mx import mlx_tests class TestAutograd(mlx_tests.MLXTestCase): def test_jvp(self): fun = lambda x: 2 * x out, dout = mx.jvp(fun, [mx.array(1.0)], [mx.array(2.0)]) self.assertEqual(out[0].item(), 2.0) self.assertEqual(dout[0].item(), 4.0) fun = lambda x, y: x * y _, out = mx.jvp( fun, [mx.array(4.0), mx.array(2.0)], [mx.array(3.0), mx.array(2.0)] ) self.assertEqual(out[0].item(), 4.0 * 2.0 + 2.0 * 3.0) fun = lambda x, y, z: (x * y, y * z) _, out = mx.jvp( fun, [mx.array(2.0), mx.array(4.0), mx.array(6.0)], [mx.array(1.0), mx.array(3.0), mx.array(1.0)], ) self.assertEqual(len(out), 2) self.assertEqual(out[0].item(), 4.0 * 1.0 + 2.0 * 3.0) self.assertEqual(out[1].item(), 4.0 * 1.0 + 6.0 * 3.0) def test_vjp(self): fun = lambda x: 2 * x out, dout = mx.vjp(fun, [mx.array(1.0)], [mx.array(2.0)]) self.assertEqual(out[0].item(), 2.0) self.assertEqual(dout[0].item(), 4.0) fun = lambda x, y: x * y _, dout = mx.vjp(fun, [mx.array(4.0), mx.array(2.0)], [mx.array(3.0)]) self.assertEqual(dout[0].item(), 6.0) self.assertEqual(dout[1].item(), 12.0) fun = lambda x, y, z: (x * y, y * z) _, out = mx.vjp( fun, [mx.array(2.0), mx.array(4.0), mx.array(6.0)], [mx.array(1.0), mx.array(3.0)], ) self.assertEqual(len(out), 3) self.assertEqual(out[0].item(), 4.0 * 1.0) self.assertEqual(out[1].item(), 2.0 * 1.0 + 6.0 * 3.0) self.assertEqual(out[2].item(), 4.0 * 3.0) def test_grad(self): fun = lambda x: x * x value, dfdx = mx.value_and_grad(fun)(mx.array(0.5)) self.assertEqual(value.item(), 0.25) self.assertEqual(dfdx.item(), 1.0) dfdx = mx.grad(fun)(mx.array(0.5)) self.assertEqual(dfdx.item(), 1.0) df2dx2 = mx.grad(mx.grad(fun))(mx.array(0.5)) self.assertEqual(df2dx2.item(), 2.0) df3dx3 = mx.grad(mx.grad(mx.grad(fun)))(mx.array(0.5)) self.assertEqual(df3dx3.item(), 0.0) fun = lambda x, y: x * y x = mx.array(2.0) y = mx.array(3.0) dfdx = mx.grad(fun, argnums=0)(x, y) self.assertEqual(dfdx.item(), 3.0) dfdx = mx.grad(fun, argnums=1)(x, y) self.assertEqual(dfdx.item(), 2.0) # Pass non array args to functions works fun = lambda x, y: x value, dfdx = mx.value_and_grad(fun)(mx.array(2.0), "hello") self.assertEqual(value.item(), 2.0) self.assertEqual(dfdx.item(), 1.0) dfdx = mx.grad(fun)(mx.array(2.0), "hello") self.assertEqual(dfdx.item(), 1.0) # Raises when function does not return array fun = lambda x: "hello" with self.assertRaises(ValueError): mx.grad(fun)(mx.array(2.0)) # Raises for invalid argument number or argument type fun = lambda x: x with self.assertRaises(ValueError): mx.grad(fun, argnums=2)(mx.array(2.0)) with self.assertRaises(ValueError): mx.grad(fun, argnums=-2)(mx.array(2.0)) with self.assertRaises(ValueError): mx.grad(fun)("hello") # Raises when output is not a scalar array fun = lambda x: mx.sum(x, keepdims=True) with self.assertRaises(ValueError): mx.grad(fun)(mx.ones((2, 2))) def test_grad_trees(self): fun = lambda x, y: x * y value, dfdx = mx.value_and_grad(fun, (0, 1))(mx.array(0.5), mx.array(2.0)) self.assertEqual(value.item(), 1.0) self.assertTrue(isinstance(dfdx, tuple)) self.assertEqual(dfdx[0].item(), 2.0) self.assertEqual(dfdx[1].item(), 0.5) fun = lambda x, y: x * y value, dfdx = mx.value_and_grad(fun, 1)(mx.array(0.5), mx.array(2.0)) self.assertEqual(value.item(), 1.0) self.assertEqual(dfdx.item(), 0.5) fun = lambda p: p["x"] * p["y"] value, dfdx = mx.value_and_grad(fun)({"x": mx.array(0.5), "y": mx.array(2.0)}) self.assertEqual(value.item(), 1.0) self.assertEqual(dfdx["x"].item(), 2.0) self.assertEqual(dfdx["y"].item(), 0.5) fun = lambda p: p["x"] * p["y"] with self.assertRaises(ValueError): mx.value_and_grad(fun)({"x": 0.5, "y": mx.array(2.0)}) with self.assertRaises(ValueError): mx.value_and_grad(fun, (0, 1))({"x": mx.array(0.5), "y": mx.array(2.0)}) fun = lambda p, b: mx.square(p[0]["foo"][2]) * b value, dfdx = mx.value_and_grad(fun)( [{"foo": [[], [], mx.array(2.0)]}], mx.array(0.5) ) self.assertEqual(value.item(), 2.0) self.assertEqual(dfdx[0]["foo"][2].item(), 2.0) fun = lambda x: x with self.assertRaises(TypeError): mx.value_and_grad(fun, (None, None)) with self.assertRaises(ValueError): mx.value_and_grad(fun, tuple()) with self.assertRaises(ValueError): mx.grad(fun, argnums=(0, 0)) def test_auxiliary_values(self): def fun(x, y): l = (x * y).sum() extra = {"loss": l, "foo": y.square() + x.square(), "bar": [1, 2, 3, y, x]} return l, extra fun_value_grad = mx.value_and_grad(fun) fun_grad = mx.grad(fun) (loss, a), b = fun_value_grad(mx.ones((2, 2)), mx.ones((2, 2))) self.assertEqual(a["loss"].item(), 4) self.assertTrue(mx.array_equal(b, mx.ones((2, 2)))) self.assertTrue(mx.array_equal(a["foo"], 2 * mx.ones((2, 2)))) self.assertEqual(a["bar"][:3], [1, 2, 3]) self.assertTrue(mx.array_equal(a["bar"][3], mx.ones((2, 2)))) self.assertTrue(mx.array_equal(a["bar"][4], mx.ones((2, 2)))) with self.assertRaises(ValueError): _ = fun_grad(mx.ones((2, 2)), mx.ones((2, 2))) def test_grad_kwargs(self): fun = lambda x, y: x * y a, b = mx.array(0.5), mx.array(2.0) dfdx = mx.grad(fun) self.assertEqual(dfdx(a, b).item(), 2.0) self.assertEqual(dfdx(a, y=b).item(), 2.0) with self.assertRaises(ValueError): dfdx(x=a, y=b).item() dfdy = mx.grad(fun, argnums=[], argnames=["y"]) with self.assertRaises(ValueError): dfdy(a, b) grads = dfdy(a, y=b) self.assertTrue(isinstance(grads, tuple)) self.assertTrue(grads[0] is None) self.assertTrue(isinstance(grads[1], dict)) self.assertEqual(grads[1]["y"].item(), 0.5) grads = dfdy(x=a, y=b) self.assertEqual(grads[1]["y"].item(), 0.5) self.assertEqual(len(grads[1]), 1) dfdxy = mx.grad(fun, argnums=[0], argnames=["y"]) with self.assertRaises(ValueError): dfdxy(a, b) with self.assertRaises(ValueError): dfdxy(x=a, y=b) grads = dfdxy(a, y=b) self.assertTrue(isinstance(grads, tuple)) self.assertEqual(grads[0].item(), 2.0) self.assertTrue(isinstance(grads[1], dict)) self.assertEqual(grads[1]["y"].item(), 0.5) fun = lambda x, y, z: x * y * z dfdxyz = mx.grad(fun, argnums=[0, 1], argnames=["z"]) c = mx.array(4.0) grads = dfdxyz(a, b, z=c) self.assertTrue(isinstance(grads, tuple)) self.assertTrue(isinstance(grads[0], tuple)) self.assertEqual(grads[0][0].item(), 8.0) self.assertEqual(grads[0][1].item(), 2.0) self.assertTrue(isinstance(grads[1], dict)) self.assertEqual(grads[1]["z"].item(), 1.0) fun = lambda x, y: x * y dfdy = mx.grad(fun, argnames=["y"]) grads = dfdy(a, y=b) self.assertTrue(isinstance(grads, tuple)) self.assertTrue(grads[0] is None) self.assertTrue(isinstance(grads[1], dict)) self.assertEqual(grads[1]["y"].item(), 0.5) def test_captured(self): a = mx.array(5.0) f = lambda x: a + x g = lambda x: a + a h = lambda x: x + x dfdx = mx.grad(f) self.assertEqual(dfdx(a).item(), 1.0) dgdx = mx.grad(g) self.assertEqual(dgdx(a).item(), 0.0) dhdx = mx.grad(h) self.assertEqual(dhdx(a).item(), 2.0) d2fdx2 = mx.grad(dfdx) self.assertEqual(d2fdx2(a).item(), 0.0) d2gdx2 = mx.grad(dgdx) self.assertEqual(d2gdx2(a).item(), 0.0) d2hdx2 = mx.grad(dhdx) self.assertEqual(d2hdx2(a).item(), 0.0) def test_stop_gradient(self): shape_in = (4, 4) w_in = mx.ones(shape_in) x_in = mx.ones(shape_in) cotan = mx.ones(shape_in) def h(w, x): x1 = 2 * x y = mx.stop_gradient(x1) y1 = 3 * y return w @ y1 vals, vjps = mx.vjp(h, [w_in, x_in], [cotan]) mx.eval(vjps) self.assertTrue(mx.allclose(vjps[0], 24.0 * mx.ones(shape_in))) self.assertTrue(mx.allclose(vjps[1], mx.zeros(shape_in))) g = lambda x: h(w_in, x) vals, vjps = mx.vjp(g, [x_in], [cotan]) mx.eval(vjps) self.assertTrue(mx.allclose(vjps[0], mx.zeros(shape_in))) def test_update_state(self): y = mx.array([1.0]) state = mx.zeros((2,)) def fn(y, x): nonlocal state x = y * x state = state + x return x.sum() x = mx.ones((2,)) mx.grad(fn)(y, x) mx.eval(state) self.assertTrue(mx.allclose(state, mx.ones((2,)))) def test_scatter_vjp(self): def fun(x, idx): x[idx] = 2.0 return x.sum() dfdx = mx.grad(fun)(mx.array([1.0, 2.0, 3.0]), mx.array([1])) self.assertTrue(mx.array_equal(dfdx, mx.array([1.0, 0.0, 1.0]))) self.assertEqual(dfdx.dtype, mx.float32) y = mx.array([0.0, 1.0, 2.0]) def fun(x, idx): y[idx] = x return y.sum() dfdx = mx.grad(fun)(mx.array([2.0]), mx.array([1])) self.assertTrue(mx.array_equal(dfdx, mx.array([1.0]))) self.assertEqual(dfdx.dtype, mx.float32) def test_scatter_max_vjp(self): def fun(src, updates): x = src.at[1].maximum(updates) return x cotan = mx.array([4.0, 5.0, 6.0]) _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0]), mx.array([[3.0]])], [cotan]) mx.eval(vjps) # Update larger than value self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 0.0, 6.0]))) self.assertTrue(mx.allclose(vjps[1], mx.array([5.0]))) cotan = mx.array([[4.0], [5.0], [6.0]]) _, vjps = mx.vjp( fun, [mx.array([[1.0], [2.0], [3.0]]), mx.array([[[2.0]]])], [cotan] ) mx.eval(vjps) # Update and value are equal self.assertTrue(mx.allclose(vjps[0], mx.array([[4.0], [5.0], [6.0]]))) self.assertTrue(mx.allclose(vjps[1], mx.array([[[5.0]]]))) def test_scatter_min_vjp(self): def fun(src, updates): x = src.at[1].minimum(updates) return x cotan = mx.array([4.0, 5.0, 6.0]) _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0]), mx.array([[3.0]])], [cotan]) mx.eval(vjps) # Update larger than value self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0]))) self.assertTrue(mx.allclose(vjps[1], mx.array([0.0]))) cotan = mx.array([[4.0], [5.0], [6.0]]) _, vjps = mx.vjp( fun, [mx.array([[1.0], [2.0], [3.0]]), mx.array([[[2.0]]])], [cotan] ) mx.eval(vjps) # Update and value are equal self.assertTrue(mx.allclose(vjps[0], mx.array([[4.0], [5.0], [6.0]]))) self.assertTrue(mx.allclose(vjps[1], mx.array([[[5.0]]]))) def test_split_against_slice(self): def f_split(x): a, _, b = x.split(3, -1) return (a * b).sum() def f_slice(x): step = x.shape[-1] // 3 a = x[..., :step] b = x[..., -step:] return (a * b).sum() x = mx.random.uniform(shape=(100, 300)) mx.eval(x) df1 = mx.grad(f_split) df2 = mx.grad(f_slice) self.assertTrue(mx.allclose(df1(x), df2(x))) def test_vjp_types(self): def fun(x): return x for t in [mx.float16, mx.bfloat16, mx.float32]: out = mx.grad(fun)(mx.array(1.0, t)) self.assertEqual(out.dtype, t) def fun(x): return x.sum() for t in [mx.float16, mx.bfloat16, mx.float32]: out = mx.grad(fun)(mx.array(1.0, t)) self.assertEqual(out.dtype, t) def fun(x, y): return (x + y).sum() for t in [mx.float16, mx.bfloat16, mx.float32]: out = mx.grad(fun)(mx.array(1.0, t), mx.array(1.0, t)) self.assertEqual(out.dtype, t) def test_power_grad(self): x = mx.array(0.0) g = mx.grad(lambda x: x**2)(x) self.assertEqual(g.item(), 0.0) x = mx.array(0.0) g = mx.grad(lambda x: x**1.5)(x) self.assertEqual(g.item(), 0.0) x = mx.array(2.0) g = mx.grad(lambda x: x**2)(x) self.assertAlmostEqual(g.item(), 4.0) def test_eval_in_grad(self): arr = mx.array([1.0]) cotan = mx.array([1.0, 1.0]) y = mx.array([2.0, 2.0]) def func(x): x = x + y cond = x < 1 cond.tolist() return x**2 _, vjps = mx.vjp(func, (arr,), (cotan,)) self.assertEqual(vjps[0].item(), 12.0) def func(x): x = x + mx.array([1.0, 1.0]) mx.eval(x) return x**2 _, vjps = mx.vjp(func, (arr,), (cotan,)) self.assertEqual(vjps[0].item(), 8.0) def test_power_grad(self): def fun(x, y): res = x - y return res**x grad = mx.grad(fun)(mx.array(1.0), mx.array(1.0)) self.assertEqual(grad.item(), 1.0) def test_cumprod_grad(self): def fun(y): return mx.cumprod(y).sum() y = mx.array([2.0, 1.0, 2.0, 2.0, 3.0]) out = mx.grad(fun)(y) expected = mx.array([20.0, 38.0, 18.0, 16.0, 8.0]) self.assertTrue(mx.allclose(out, expected)) y = mx.array([2.0, 0.0, 2.0, 2.0, 3.0]) out = mx.grad(fun)(y) expected = mx.array([1.0, 38.0, 0.0, 0.0, 0.0]) self.assertTrue(mx.allclose(out, expected)) y = mx.array([2.0, 0.0, 2.0, 0.0, 3.0]) out = mx.grad(fun)(y) expected = mx.array([1.0, 6.0, 0.0, 0.0, 0.0]) self.assertTrue(mx.allclose(out, expected)) def fun(y): return mx.cumprod(y, inclusive=False).sum() y = mx.array([2.0, 1.0, 2.0, 2.0, 3.0]) out = mx.grad(fun)(y) expected = mx.array([8.0, 14.0, 6.0, 4.0, 0.0]) self.assertTrue(mx.allclose(out, expected)) y = mx.array([2.0, 0.0, 2.0, 2.0, 3.0]) out = mx.grad(fun)(y) expected = mx.array([1.0, 14.0, 0.0, 0.0, 0.0]) self.assertTrue(mx.allclose(out, expected)) y = mx.array([2.0, 0.0, 2.0, 0.0, 3.0]) out = mx.grad(fun)(y) expected = mx.array([1.0, 6.0, 0.0, 0.0, 0.0]) self.assertTrue(mx.allclose(out, expected)) def fun(y): return mx.cumprod(y, inclusive=False, reverse=True).sum() y = mx.array([2.0, 1.0, 2.0, 2.0, 3.0]) out = mx.grad(fun)(y) expected = mx.array([0.0, 12.0, 12.0, 15.0, 11.0]) self.assertTrue(mx.allclose(out, expected)) y = mx.array([2.0, 0.0, 2.0, 2.0, 3.0]) out = mx.grad(fun)(y) expected = mx.array([0.0, 12.0, 6.0, 9.0, 7.0]) self.assertTrue(mx.allclose(out, expected)) y = mx.array([2.0, 0.0, 2.0, 0.0, 3.0]) out = mx.grad(fun)(y) expected = mx.array([0.0, 0.0, 0.0, 9.0, 1.0]) self.assertTrue(mx.allclose(out, expected)) def fun(y): return mx.cumprod(y, reverse=True).sum() y = mx.array([2.0, 1.0, 2.0, 2.0, 3.0]) out = mx.grad(fun)(y) expected = mx.array([12.0, 36.0, 24.0, 27.0, 19.0]) self.assertTrue(mx.allclose(out, expected)) y = mx.array([2.0, 0.0, 2.0, 2.0, 3.0]) out = mx.grad(fun)(y) expected = mx.array([0.0, 36.0, 6.0, 9.0, 7.0]) self.assertTrue(mx.allclose(out, expected)) y = mx.array([2.0, 0.0, 2.0, 0.0, 3.0]) out = mx.grad(fun)(y) expected = mx.array([0.0, 0.0, 0.0, 9.0, 1.0]) self.assertTrue(mx.allclose(out, expected)) def test_topk_grad(self): a = mx.array([[1, 2, 6, 4, 5], [9, 5, 6, 7, 8]], mx.float32) def fun(x): return mx.topk(x, 2) out = mx.vjp(fun, (a,), (mx.ones((2, 2)),))[1][0] expected = mx.array([[0, 0, 1, 0, 1], [1, 0, 0, 0, 1]], mx.float32) self.assertTrue(mx.array_equal(out, expected)) def test_custom_function(self): # Make a custom function my_exp = mx.custom_function(mx.exp) # Ensure everything works dy = mx.grad(my_exp)(mx.array(1.0)) self.assertTrue(mx.allclose(dy, mx.exp(mx.array(1.0)))) (ex,), (dex,) = mx.jvp(my_exp, [mx.array(1.0)], [mx.array(1.0)]) self.assertTrue(mx.allclose(dex, mx.exp(mx.array(1.0)))) self.assertTrue(mx.allclose(ex, dex)) ex = mx.vmap(my_exp)(mx.ones(10)) self.assertTrue(mx.allclose(ex, mx.exp(mx.ones(10)))) # Ensure that the vjp is being overriden but everything else still # works. @my_exp.vjp def my_exp_vjp(x, dx, ex): return mx.ones_like(x) * 42 dy = mx.grad(my_exp)(mx.array(1.0)) self.assertTrue(mx.allclose(dy, mx.array(42.0))) (ex,), (dex,) = mx.jvp(my_exp, [mx.array(1.0)], [mx.array(1.0)]) self.assertTrue(mx.allclose(dex, mx.exp(mx.array(1.0)))) self.assertTrue(mx.allclose(ex, dex)) ex = mx.vmap(my_exp)(mx.ones(10)) self.assertTrue(mx.allclose(ex, mx.exp(mx.ones(10)))) # Ensure that setting the jvp and vmap also works. @my_exp.jvp def my_exp_jvp(x, dx): return mx.ones_like(x) * 7 * dx @my_exp.vmap def my_exp_vmap(x, axis): return mx.ones_like(x) * 3, axis dy = mx.grad(my_exp)(mx.array(1.0)) self.assertTrue(mx.allclose(dy, mx.array(42.0))) (ex,), (dex,) = mx.jvp(my_exp, [mx.array(1.0)], [mx.array(1.0)]) self.assertTrue(mx.allclose(dex, mx.array(7.0))) self.assertTrue(mx.allclose(ex, mx.exp(mx.array(1.0)))) ex = mx.vmap(my_exp)(mx.ones(10)) self.assertTrue(mx.allclose(ex, 3 * mx.ones(10))) # Test pytrees @mx.custom_function def my_double(params): return {"out": 2 * params["x"] * params["y"]} dy = mx.grad(lambda p: my_double(p)["out"].sum())( {"x": mx.ones(2), "y": mx.ones(2)} ) self.assertTrue(mx.allclose(dy["x"], mx.ones(2) * 2)) self.assertTrue(mx.allclose(dy["y"], mx.ones(2) * 2)) @my_double.vjp def random_grads(primals, cotangents, outputs): return {"x": mx.zeros_like(primals["x"]), "y": mx.ones_like(primals["y"])} dy = mx.grad(lambda p: my_double(p)["out"].sum())( {"x": mx.ones(2), "y": mx.ones(2)} ) self.assertTrue(mx.allclose(dy["x"], mx.zeros(2))) self.assertTrue(mx.allclose(dy["y"], mx.ones(2))) def outer_f(a, b): return my_double({"x": a, "y": b})["out"] inputs = [mx.random.normal(shape=(2,)) for i in range(2)] tans = [mx.random.normal(shape=(2,)) for i in range(2)] out1, dout1 = mx.jvp(outer_f, inputs, tans) @my_double.jvp def random_grads(primals, tangents): return { "out": 2 * primals["x"] * tangents["y"] + 2 * primals["y"] * tangents["x"] + 1 } out2, dout2 = mx.jvp(outer_f, inputs, tans) self.assertTrue(mx.allclose(out1[0], out2[0])) self.assertTrue(mx.allclose(dout1[0] + 1, dout2[0])) def test_complex_vjps(self): def fun(x): return (2.0 * mx.real(x)).sum() x = mx.array([0.0 + 1j, 1.0 + 0.0j, 0.5 + 0.5j]) dfdx = mx.grad(fun)(x) self.assertTrue(mx.allclose(dfdx, 2 * mx.ones_like(x))) def fun(x): return (2.0 * mx.imag(x)).sum() x = mx.array([0.0 + 1j, 1.0 + 0.0j, 0.5 + 0.5j]) dfdx = mx.grad(fun)(x) self.assertTrue(mx.allclose(dfdx, -2j * mx.ones_like(x))) def test_flatten_unflatten_vjps(self): def fun(x): y = mx.unflatten(x, 0, (2, 2)) return y.sum() x = mx.zeros((4, 8)) self.assertEqual(mx.grad(fun)(x).shape, (4, 8)) def fun(x): y = mx.flatten(x, 0, 2) return y.sum() x = mx.zeros((2, 4, 8)) self.assertEqual(mx.grad(fun)(x).shape, (2, 4, 8)) def test_concatenate_vjps(self): def fun(x, y): return mx.concatenate([x, y]) x = mx.array([1, 2, 3], mx.float32) y = mx.array([1, 2, 3], mx.float16) grads = mx.vjp(fun, (x, y), (mx.ones((6,)),))[1] self.assertTrue(mx.allclose(grads[0], mx.ones(3))) self.assertTrue(mx.allclose(grads[1], mx.ones(3))) self.assertEqual(grads[0].dtype, mx.float32) self.assertEqual(grads[1].dtype, mx.float16) def test_matmul_jvps(self): a = mx.random.uniform(shape=(4, 4)) b = mx.random.uniform(shape=(4, 4)) c = mx.random.uniform(shape=(4, 4)) d = mx.random.uniform(shape=(4, 4)) _, tangent = mx.jvp(lambda a: a @ b, (a,), (c,)) self.assertTrue(mx.allclose(tangent[0], c @ b)) _, tangent = mx.jvp(lambda b: a @ b, (b,), (d,)) self.assertTrue(mx.allclose(tangent[0], a @ d)) _, tangent = mx.jvp(lambda a, b: a @ b, (a, b), (c, d)) self.assertTrue(mx.allclose(tangent[0], a @ d + c @ b)) x = mx.random.uniform(shape=(4, 4)) y = mx.random.uniform(shape=(4, 4)) z = mx.random.uniform(shape=(4, 4)) _, (tangent,) = mx.jvp(lambda a, b, c: a @ b + c, (a, b, c), (x, y, z)) _, (expected,) = mx.jvp(lambda a, b, c: mx.addmm(c, a, b), (a, b, c), (x, y, z)) self.assertTrue(mx.allclose(tangent, expected)) _, (tangent,) = mx.jvp(lambda a, c: a @ b + c, (a, c), (x, z)) _, (expected,) = mx.jvp(lambda a, c: mx.addmm(c, a, b), (a, c), (x, z)) self.assertTrue(mx.allclose(tangent, expected)) _, (tangent,) = mx.jvp(lambda b, c: a @ b + c, (b, c), (y, z)) _, (expected,) = mx.jvp(lambda b, c: mx.addmm(c, a, b), (b, c), (y, z)) self.assertTrue(mx.allclose(tangent, expected)) _, (tangent,) = mx.jvp(lambda c: a @ b + c, (c,), (z,)) _, (expected,) = mx.jvp(lambda c: mx.addmm(c, a, b), (c,), (z,)) self.assertTrue(mx.allclose(tangent, expected)) def test_put_along_axis_grads(self): a = mx.zeros((5, 1)) b = mx.ones((2, 1)) def fun(a, b): idx = mx.array([[0], [3]]) return mx.put_along_axis(a, idx, b, axis=0) # Test VJP cotan = mx.full((5, 1), 2.0) _, (da, db) = mx.vjp(fun, (a, b), (cotan,)) expected_da = mx.array([0.0, 2.0, 2.0, 0.0, 2.0])[:, None] expected_db = mx.array([2.0, 2.0])[:, None] self.assertTrue(mx.allclose(expected_da, da)) self.assertTrue(mx.allclose(expected_db, db)) # Test JVP tan_a = mx.full((5, 1), 2.0) tan_b = mx.full((2, 1), 3.0) _, (jout,) = mx.jvp(fun, (a, b), (tan_a, tan_b)) expected = mx.array([3.0, 2.0, 2.0, 3.0, 2.0])[:, None] self.assertTrue(mx.allclose(expected, jout)) def fun(a): idx = mx.array([[0], [3]]) return mx.put_along_axis(a, idx, b, axis=0) _, (jout,) = mx.jvp(fun, (a,), (tan_a,)) expected = mx.array([0.0, 2.0, 2.0, 0.0, 2.0])[:, None] self.assertTrue(mx.allclose(expected, jout)) if __name__ == "__main__": unittest.main()