MCPcopy Create free account
hub / github.com/ml-explore/mlx / test_stack

Method test_stack

python/tests/test_ops.py:2378–2407  ·  view source on GitHub ↗
(self)

Source from the content-addressed store, hash-verified

2376 self.assertCmpNumpy([5, 6], mx.eye, np.eye, k=-2)
2377
2378 def test_stack(self):
2379 a = mx.ones((2,))
2380 np_a = np.ones((2,))
2381 b = mx.ones((2,))
2382 np_b = np.ones((2,))
2383
2384 # One dimensional stack axis=0
2385 c = mx.stack([a, b])
2386 np_c = np.stack([np_a, np_b])
2387 self.assertTrue(np.array_equal(c, np_c))
2388
2389 # One dimensional stack axis=1
2390 c = mx.stack([a, b], axis=1)
2391 np_c = np.stack([np_a, np_b], axis=1)
2392 self.assertTrue(np.array_equal(c, np_c))
2393
2394 a = mx.ones((1, 2))
2395 np_a = np.ones((1, 2))
2396 b = mx.ones((1, 2))
2397 np_b = np.ones((1, 2))
2398
2399 # Two dimensional stack axis=0
2400 c = mx.stack([a, b])
2401 np_c = np.stack([np_a, np_b])
2402 self.assertTrue(np.array_equal(c, np_c))
2403
2404 # Two dimensional stack axis=1
2405 c = mx.stack([a, b], axis=1)
2406 np_c = np.stack([np_a, np_b], axis=1)
2407 self.assertTrue(np.array_equal(c, np_c))
2408
2409 def test_flatten(self):
2410 x = mx.zeros([2, 3, 4])

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected