def test_squeeze(self): torch.manual_seed(1) data = ht.random.randn(1, 4, 5, 1, device=ht_device) # 4D local tensor, no axis result = ht.squeeze(data) self.assertIsInstance(result, ht.DNDarray) self.assertEqual(result.dtype, ht.float64) self.assertEqual(result._DNDarray__array.dtype, torch.float64) self.assertEqual(result.shape, (4, 5)) self.assertEqual(result.lshape, (4, 5)) self.assertEqual(result.split, None) self.assertTrue((result._DNDarray__array == data._DNDarray__array.squeeze()).all()) # 4D local tensor, major axis result = ht.squeeze(data, axis=0) self.assertIsInstance(result, ht.DNDarray) self.assertEqual(result.dtype, ht.float64) self.assertEqual(result._DNDarray__array.dtype, torch.float64) self.assertEqual(result.shape, (4, 5, 1)) self.assertEqual(result.lshape, (4, 5, 1)) self.assertEqual(result.split, None) self.assertTrue((result._DNDarray__array == data._DNDarray__array.squeeze(0)).all()) # 4D local tensor, minor axis result = ht.squeeze(data, axis=-1) self.assertIsInstance(result, ht.DNDarray) self.assertEqual(result.dtype, ht.float64) self.assertEqual(result._DNDarray__array.dtype, torch.float64) self.assertEqual(result.shape, (1, 4, 5)) self.assertEqual(result.lshape, (1, 4, 5)) self.assertEqual(result.split, None) self.assertTrue((result._DNDarray__array == data._DNDarray__array.squeeze(-1)).all()) # 4D local tensor, tuple axis result = data.squeeze(axis=(0, -1)) self.assertIsInstance(result, ht.DNDarray) self.assertEqual(result.dtype, ht.float64) self.assertEqual(result._DNDarray__array.dtype, torch.float64) self.assertEqual(result.shape, (4, 5)) self.assertEqual(result.lshape, (4, 5)) self.assertEqual(result.split, None) self.assertTrue((result._DNDarray__array == data._DNDarray__array.squeeze()).all()) # 4D split tensor, along the axis # TODO: reinstate this test of uneven dimensions distribution # after update to Allgatherv implementation (Issue #273 depending on #233) # data = ht.array(ht.random.randn(1, 4, 5, 1), split=1) # result = ht.squeeze(data, axis=-1) # self.assertIsInstance(result, ht.DNDarray) # # TODO: the following works locally but not when distributed, # #self.assertEqual(result.dtype, ht.float32) # #self.assertEqual(result._DNDarray__array.dtype, torch.float32) # self.assertEqual(result.shape, (1, 12, 5)) # self.assertEqual(result.lshape, (1, 12, 5)) # self.assertEqual(result.split, 1) # 3D split tensor, across the axis size = ht.MPI_WORLD.size * 2 data = ht.triu(ht.ones((1, size, size), split=1, device=ht_device), k=1) result = ht.squeeze(data, axis=0) self.assertIsInstance(result, ht.DNDarray) # TODO: the following works locally but not when distributed, # self.assertEqual(result.dtype, ht.float32) # self.assertEqual(result._DNDarray__array.dtype, torch.float32) self.assertEqual(result.shape, (size, size)) self.assertEqual(result.lshape, (size, size)) # self.assertEqual(result.split, None) # check exceptions with self.assertRaises(ValueError): data.squeeze(axis=(0, 1)) with self.assertRaises(TypeError): data.squeeze(axis=1.1) with self.assertRaises(TypeError): data.squeeze(axis="y") with self.assertRaises(ValueError): ht.argmin(data, axis=-4)
def logsumexp(self, a, axis=None, b=None, keepdim=False, return_sign=False): """ Adapted to HeAT from scikit-learn. Compute the log of the sum of exponentials of input elements. Parameters ---------- a : ht.tensor Input array. axis : None or int or tuple of ints, optional Axis or axes over which the sum is taken. By default `axis` is None, and all elements are summed. keepdim : bool, optional If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the original array. b : ht.tensor, optional Scaling factor for exp(`a`) must be of the same shape as `a` or broadcastable to `a`. These values may be negative in order to implement subtraction. #return_sign : bool, optional If this is set to True, the result will be a pair containing sign information; if False, results that are negative will be returned as NaN. Default is False (no sign information). #TODO: returns NotImplementedYet error. Returns ------- res : ht.tensor The result, ``np.log(np.sum(np.exp(a)))`` calculated in a numerically more stable way. If `b` is given then ``np.log(np.sum(b*np.exp(a)))`` is returned. #TODO sgn : ndarray NOT IMPLEMENTED YET If return_sign is True, this will be an array of floating-point numbers matching res and +1, 0, or -1 depending on the sign of the result. If False, only one result is returned. """ if b is not None: raise NotImplementedError("Not implemented for weighted logsumexp") a_max = ht.max(a, axis=axis, keepdim=True) # TODO: sanitize a_max / implement isfinite(): sanitation module, cf. #468 # if a_max.numdims > 0: # a_max[~np.isfinite(a_max)] = 0 # elif not np.isfinite(a_max): # a_max = 0 # TODO: reinstate after allowing b not None # if b is not None: # b = np.asarray(b) # tmp = b * np.exp(a - a_max) # else: tmp = ht.exp(a - a_max) s = ht.sum(tmp, axis=axis, keepdim=keepdim) if return_sign: raise NotImplementedError("Not implemented for return_sign") # sgn = np.sign(s) # TODO: np.sign # s *= sgn # /= makes more sense but we need zero -> zero out = ht.log(s) if not keepdim: a_max = ht.squeeze(a_max, axis=axis) out += a_max # if return_sign: #TODO: np.sign # return out, sgn # else: return out