def test_masked_fill_(self):
t = TensorBase(np.ones((2, 3)))
value = 2.0
mask = TensorBase([[0, 0, 0], [1, 1, 0]])
t.masked_fill_(mask, value)
self.assertTrue(
np.array_equal(t, TensorBase([[1.0, 1.0, 1.0], [2.0, 2.0, 1.0]])))
def testMasked_fill_broadcasting(self):
t = TensorBase(np.ones((2, 3)))
value = 2
mask = TensorBase([[1], [0]])
t.masked_fill_(mask, value)
self.assertTrue(np.array_equal(t, TensorBase([[2, 2, 2], [1, 1, 1]])))
