| 235 | ) |
| 236 | |
| 237 | def test_bernoulli(self): |
| 238 | a = mx.random.bernoulli() |
| 239 | self.assertEqual(a.shape, ()) |
| 240 | self.assertEqual(a.dtype, mx.bool_) |
| 241 | |
| 242 | a = mx.random.bernoulli(mx.array(0.5), [5]) |
| 243 | self.assertEqual(a.shape, (5,)) |
| 244 | |
| 245 | a = mx.random.bernoulli(mx.array([2.0, -2.0])) |
| 246 | self.assertEqual(a.tolist(), [True, False]) |
| 247 | self.assertEqual(a.shape, (2,)) |
| 248 | |
| 249 | p = mx.array([0.1, 0.2, 0.3]) |
| 250 | mx.reshape(p, [1, 3]) |
| 251 | x = mx.random.bernoulli(p, [4, 3]) |
| 252 | self.assertEqual(x.shape, (4, 3)) |
| 253 | |
| 254 | with self.assertRaises(ValueError): |
| 255 | mx.random.bernoulli(p, [2]) # Bad shape |
| 256 | |
| 257 | with self.assertRaises(ValueError): |
| 258 | mx.random.bernoulli(0, [2]) # Bad type |
| 259 | |
| 260 | def test_truncated_normal(self): |
| 261 | a = mx.random.truncated_normal(-2.0, 2.0) |