def test_stack(): # vector shape = (100,) py_rand.set_seed() py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) py_res = py_matrix.stack([py_mat, py_mat], axis=0) torch_res = torch_matrix.stack([torch_mat, torch_mat], axis=0) assert_close(py_res, torch_res, "stack: vectors, axis=0") py_res = py_matrix.stack([py_mat, py_mat], axis=1) torch_res = torch_matrix.stack([torch_mat, torch_mat], axis=1) assert_close(py_res, torch_res, "stack: vectors, axis=1") # matrix shape = (100, 100) py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) py_res = py_matrix.stack([py_mat, py_mat], axis=0) torch_res = torch_matrix.stack([torch_mat, torch_mat], axis=0) assert_close(py_res, torch_res, "stack: matrices, axis=0") py_res = py_matrix.stack([py_mat, py_mat], axis=1) torch_res = torch_matrix.stack([torch_mat, torch_mat], axis=1) assert_close(py_res, torch_res, "stack: matrices, axis=1")
def test_clip_inplace(): shape = (100,100) py_rand.set_seed() py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) # test two sided clip py_matrix.clip_inplace(py_mat, a_min=0, a_max=1) torch_matrix.clip_inplace(torch_mat, a_min=0, a_max=1) assert_close(py_mat, torch_mat, "clip_inplace (two-sided)") # test lower clip py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) py_matrix.clip_inplace(py_mat, a_min=0) torch_matrix.clip_inplace(torch_mat, a_min=0) assert_close(py_mat, torch_mat, "clip_inplace (lower)") # test upper clip py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) py_matrix.clip_inplace(py_mat, a_max=1) torch_matrix.clip_inplace(torch_mat, a_max=1) assert_close(py_mat, torch_mat, "clip_inplace (upper)")
def test_fast_energy_distance(): shape = (1000, 50) downsample = 100 # close distributions a_mean, a_scale = 0, 1 b_mean, b_scale = 0.1, 0.9 py_rand.set_seed() py_a = a_mean + a_scale * py_rand.randn(shape) py_b = b_mean + b_scale * py_rand.randn(shape) torch_a = torch_matrix.float_tensor(py_a) torch_b = torch_matrix.float_tensor(py_b) py_dist = py_matrix.fast_energy_distance(py_a, py_b, downsample=downsample) torch_dist = torch_matrix.fast_energy_distance(torch_a, torch_b, downsample=downsample) # python fast_energy_distance is stochastic even after calling # py_rand.set_seed() because it uses the numba random number # generator rather than the one in numpy. # this test can fail stochastically assert py_dist < 0.25, \ "python energy distance is too big" assert torch_dist < 0.25, \ "torch energy distance is too big" # distance distributions a_mean, a_scale = 1, 1 b_mean, b_scale = -1, 1 py_rand.set_seed() py_a = a_mean + a_scale * py_rand.randn(shape) py_b = b_mean + b_scale * py_rand.randn(shape) torch_a = torch_matrix.float_tensor(py_a) torch_b = torch_matrix.float_tensor(py_b) py_dist = py_matrix.fast_energy_distance(py_a, py_b, downsample=downsample) torch_dist = torch_matrix.fast_energy_distance(torch_a, torch_b, downsample=downsample) assert py_dist > 10, \ "python energy distance is too small" assert torch_dist > 10, \ "torch energy distance is too small"
def test_logaddexp(): shape = (100, 100) py_rand.set_seed() py_x_1 = py_rand.randn(shape) py_x_2 = py_rand.randn(shape) torch_x_1 = torch_matrix.float_tensor(py_x_1) torch_x_2 = torch_matrix.float_tensor(py_x_2) py_y = py_func.logaddexp(py_x_1, py_x_2) torch_y = torch_func.logaddexp(torch_x_1, torch_x_2) assert_close(py_y, torch_y, "logaddexp")
def test_tall(): shape = (100, 100) py_rand.set_seed() py_x = py_rand.randn(shape) py_y = py_rand.randn(shape) torch_x = torch_matrix.float_tensor(py_x) torch_y = torch_matrix.float_tensor(py_y) py_res = py_matrix.lesser_equal(py_x, py_y) torch_res = torch_matrix.lesser_equal(torch_x, torch_y) # overall py_all = py_matrix.tall(py_res) torch_all = torch_matrix.tall(torch_res) assert py_all == torch_all, \ "python tall != torch tall: overall" # axis = 0 py_all = py_matrix.tall(py_res, axis=0) torch_all = torch_matrix.tall(torch_res, axis=0) py_torch_all = torch_matrix.to_numpy_array(torch_all) assert py_matrix.allclose(py_all, py_torch_all), \ "python tall != torch tall: (axis-0)" # axis = 1 py_all = py_matrix.tall(py_res, axis=1) torch_all = torch_matrix.tall(torch_res, axis=1) py_torch_all = torch_matrix.to_numpy_array(torch_all) assert py_matrix.allclose(py_all, py_torch_all), \ "python tall != torch tall: (axis-1)" # axis = 0, keepdims py_all = py_matrix.tall(py_res, axis=0, keepdims=True) torch_all = torch_matrix.tall(torch_res, axis=0, keepdims=True) py_torch_all = torch_matrix.to_numpy_array(torch_all) assert py_matrix.allclose(py_all, py_torch_all), \ "python tall != torch tall: (axis-0, keepdims)" # axis = 1, keepdims py_all = py_matrix.tall(py_res, axis=1, keepdims=True) torch_all = torch_matrix.tall(torch_res, axis=1, keepdims=True) py_torch_all = torch_matrix.to_numpy_array(torch_all) assert py_matrix.allclose(py_all, py_torch_all), \ "python tall != torch tall: (axis-1, keepdim)"
def test_minimum(): shape = (100, 100) py_rand.set_seed() py_x = py_rand.randn(shape) py_y = py_rand.randn(shape) torch_x = torch_matrix.float_tensor(py_x) torch_y = torch_matrix.float_tensor(py_y) py_res = py_matrix.minimum(py_x, py_y) torch_res = torch_matrix.minimum(torch_x, torch_y) assert_close(py_res, torch_res, "minimum")
def test_sqrt_div(): shape = (100,100) py_rand.set_seed() py_x = py_rand.randn(shape) py_y = py_rand.randn(shape) ** 2 torch_x = torch_matrix.float_tensor(py_x) torch_y = torch_matrix.float_tensor(py_y) py_sqrt_div = py_matrix.sqrt_div(py_x, py_y) torch_sqrt_div = torch_matrix.sqrt_div(torch_x, torch_y) assert_close(py_sqrt_div, torch_sqrt_div, "sqrt_div")
def test_outer(): a_shape = (100,) b_shape = (100,) py_rand.set_seed() py_a = py_rand.randn(a_shape) py_b = py_rand.randn(b_shape) torch_a = torch_matrix.float_tensor(py_a) torch_b = torch_matrix.float_tensor(py_b) py_res = py_matrix.outer(py_a, py_b) torch_res = torch_matrix.outer(torch_a, torch_b) assert_close(py_res, torch_res, "outer")
def test_squared_euclidean_distance(): shape = (100,) py_rand.set_seed() py_a = py_rand.randn(shape) py_b = py_rand.randn(shape) torch_a = torch_matrix.float_tensor(py_a) torch_b = torch_matrix.float_tensor(py_b) py_dist = py_matrix.squared_euclidean_distance(py_a, py_b) torch_dist = torch_matrix.squared_euclidean_distance(torch_a, torch_b) assert allclose(py_dist, torch_dist), \ "squared_euclidean_distance failure"
def test_not_equal(): shape = (100, 100) py_rand.set_seed() py_x = py_rand.randn(shape) py_y = py_rand.randn(shape) torch_x = torch_matrix.float_tensor(py_x) torch_y = torch_matrix.float_tensor(py_y) py_neq = py_matrix.not_equal(py_x, py_y) torch_neq = torch_matrix.not_equal(torch_x, torch_y) py_torch_neq = torch_matrix.to_numpy_array(torch_neq) assert py_matrix.allclose(py_neq, py_torch_neq), \ "python not equal != torch not equal"
def test_tprod(): shape = (100, 100) py_rand.set_seed() py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) # overall tprod py_tprod = py_matrix.tprod(py_mat) torch_tprod = torch_matrix.tprod(torch_mat) assert allclose(py_tprod, torch_tprod), \ "python overal tprod != torch overall tprod" # tprod over axis 0 py_tprod = py_matrix.tprod(py_mat, axis=0) torch_tprod = torch_matrix.tprod(torch_mat, axis=0) assert_close(py_tprod, torch_tprod, "tprod (axis-0)") # tprod over axis 1 py_tprod = py_matrix.tprod(py_mat, axis=1) torch_tprod = torch_matrix.tprod(torch_mat, axis=1) assert_close(py_tprod, torch_tprod, "tprod (axis-1)") # tprod over axis 0, keepdims = True py_tprod = py_matrix.tprod(py_mat, axis=0, keepdims=True) torch_tprod = torch_matrix.tprod(torch_mat, axis=0) assert_close(py_tprod, torch_tprod, "tprod (axis-0, keepdims)") # tprod over axis 1, keepdims = True py_tprod = py_matrix.tprod(py_mat, axis=1, keepdims=True) torch_tprod = torch_matrix.tprod(torch_mat, axis=1, keepdims=True) assert_close(py_tprod, torch_tprod, "tprod (axis-1, keepdims)")
def test_var(): shape = (100, 100) py_rand.set_seed() py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) # overall var py_var = py_matrix.var(py_mat) torch_var = torch_matrix.var(torch_mat) assert allclose(py_var, torch_var), \ "python overal var != torch overall var" # var over axis 0 py_var = py_matrix.var(py_mat, axis=0) torch_var = torch_matrix.var(torch_mat, axis=0) assert_close(py_var, torch_var, "var (axis-0)") # var over axis 1 py_var = py_matrix.var(py_mat, axis=1) torch_var = torch_matrix.var(torch_mat, axis=1) assert_close(py_var, torch_var, "var (axis-1)") # var over axis 0, keepdims = True py_var = py_matrix.var(py_mat, axis=0, keepdims=True) torch_var = torch_matrix.var(torch_mat, axis=0) assert_close(py_var, torch_var, "var (axis-0, keepdims)") # var over axis 1, keepdims = True py_var = py_matrix.var(py_mat, axis=1, keepdims=True) torch_var = torch_matrix.var(torch_mat, axis=1, keepdims=True) assert_close(py_var, torch_var, "var (axis-1, keepdims)")
def test_tsum(): shape = (100, 100) py_rand.set_seed() py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) # overall tsum py_tsum = py_matrix.tsum(py_mat) torch_tsum = torch_matrix.tsum(torch_mat) assert allclose(py_tsum, torch_tsum), \ "python overal tsum != torch overall tsum" # tsum over axis 0 py_tsum = py_matrix.tsum(py_mat, axis=0) torch_tsum = torch_matrix.tsum(torch_mat, axis=0) assert_close(py_tsum, torch_tsum, "tsum (axis-0)") # tsum over axis 1 py_tsum = py_matrix.tsum(py_mat, axis=1) torch_tsum = torch_matrix.tsum(torch_mat, axis=1) assert_close(py_tsum, torch_tsum, "tsum (axis-1)") # tsum over axis 0, keepdims = True py_tsum = py_matrix.tsum(py_mat, axis=0, keepdims=True) torch_tsum = torch_matrix.tsum(torch_mat, axis=0) assert_close(py_tsum, torch_tsum, "tsum (axis-0, keepdims)") # tsum over axis 1, keepdims = True py_tsum = py_matrix.tsum(py_mat, axis=1, keepdims=True) torch_tsum = torch_matrix.tsum(torch_mat, axis=1, keepdims=True) assert_close(py_tsum, torch_tsum, "tsum (axis-1, keepdims)")
def test_tmax(): shape = (100, 100) py_rand.set_seed() py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) # overall max py_max = py_matrix.tmax(py_mat) torch_max = torch_matrix.tmax(torch_mat) assert allclose(py_max, torch_max), \ "python overal max != torch overall max" # max over axis 0 py_max = py_matrix.tmax(py_mat, axis=0) torch_max = torch_matrix.tmax(torch_mat, axis=0) assert_close(py_max, torch_max, "tmax (axis-0)") # max over axis 1 py_max = py_matrix.tmax(py_mat, axis=1) torch_max = torch_matrix.tmax(torch_mat, axis=1) assert_close(py_max, torch_max, "tmax (axis-1)") # max over axis 0, keepdims = True py_max = py_matrix.tmax(py_mat, axis=0, keepdims=True) torch_max = torch_matrix.tmax(torch_mat, axis=0) assert_close(py_max, torch_max, "tmax (axis-0, keepdims)") # max over axis 1, keepdims = True py_max = py_matrix.tmax(py_mat, axis=1, keepdims=True) torch_max = torch_matrix.tmax(torch_mat, axis=1, keepdims=True) assert_close(py_max, torch_max, "tmax (axis-1, keepdims)")
def test_mean(): shape = (100, 100) py_rand.set_seed() py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) # overall mean py_mean = py_matrix.mean(py_mat) torch_mean = torch_matrix.mean(torch_mat) assert allclose(py_mean, torch_mean), \ "python overal mean != torch overall mean" # mean over axis 0 py_mean = py_matrix.mean(py_mat, axis=0) torch_mean = torch_matrix.mean(torch_mat, axis=0) assert_close(py_mean, torch_mean, "mean (axis-0)") # mean over axis 1 py_mean = py_matrix.mean(py_mat, axis=1) torch_mean = torch_matrix.mean(torch_mat, axis=1) assert_close(py_mean, torch_mean, "mean (axis-1)") # mean over axis 0, keepdims = True py_mean = py_matrix.mean(py_mat, axis=0, keepdims=True) torch_mean = torch_matrix.mean(torch_mat, axis=0) assert_close(py_mean, torch_mean, "mean (axis-0, keepdims)") # mean over axis 1, keepdims = True py_mean = py_matrix.mean(py_mat, axis=1, keepdims=True) torch_mean = torch_matrix.mean(torch_mat, axis=1, keepdims=True) assert_close(py_mean, torch_mean, "mean (axis-1, keepdims)")
def test_square_mix_inplace(): shape = (100,100) torch_w = 0.1 py_w = py_matrix.float_scalar(torch_w) py_rand.set_seed() py_x = py_rand.randn(shape) py_y = py_rand.randn(shape) torch_x = torch_matrix.float_tensor(py_x) torch_y = torch_matrix.float_tensor(py_y) py_matrix.square_mix_inplace(py_w, py_x, py_y) torch_matrix.square_mix_inplace(torch_w, torch_x, torch_y) assert_close(py_x, torch_x, "square_mix_inplace")
def test_lesser_equal(): shape = (100, 100) py_rand.set_seed() py_x = py_rand.randn(shape) py_y = py_rand.randn(shape) torch_x = torch_matrix.float_tensor(py_x) torch_y = torch_matrix.float_tensor(py_y) py_res = py_matrix.lesser_equal(py_x, py_y) torch_res = torch_matrix.lesser_equal(torch_x, torch_y) py_torch_res = torch_matrix.to_numpy_array(torch_res) assert py_matrix.allclose(py_res, py_torch_res), \ "python lesser_equal != torch lesser_equal"
def test_greater(): shape = (100, 100) py_rand.set_seed() py_x = py_rand.randn(shape) py_y = py_rand.randn(shape) torch_x = torch_matrix.float_tensor(py_x) torch_y = torch_matrix.float_tensor(py_y) py_res = py_matrix.greater(py_x, py_y) torch_res = torch_matrix.greater(torch_x, torch_y) py_torch_res = torch_matrix.to_numpy_array(torch_res) assert py_matrix.allclose(py_res, py_torch_res), \ "python greater != torch greater"
def test_logcosh(): shape = (100, 100) py_rand.set_seed() py_x = py_rand.randn(shape) torch_x = torch_matrix.float_tensor(py_x) py_y = py_func.logcosh(py_x) torch_y = torch_func.logcosh(torch_x) assert_close(py_y, torch_y, "logcosh")
def test_batch_outer(): L = 10 N = 100 M = 50 v_shape = (L, N) h_shape = (L, M) py_rand.set_seed() py_v = py_rand.randn(v_shape) py_h = py_rand.randn(h_shape) torch_v = torch_matrix.float_tensor(py_v) torch_h = torch_matrix.float_tensor(py_h) py_res = py_matrix.batch_outer(py_v, py_h) torch_res = torch_matrix.batch_outer(torch_v, torch_h) assert_close(py_res, torch_res, "batch_outer")
def test_sin(): shape = (100, 100) py_rand.set_seed() py_x = py_rand.randn(shape) torch_x = torch_matrix.float_tensor(py_x) py_y = py_func.sin(py_x) torch_y = torch_func.sin(torch_x) assert_close(py_y, torch_y, "sin")
def test_sign(): shape = (100,100) py_rand.set_seed() py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) py_sign = py_matrix.sign(py_mat) torch_sign = torch_matrix.sign(torch_mat) assert_close(py_sign, torch_sign, "sign")
def test_tpow(): shape = (100, 100) power = 3 py_rand.set_seed() py_x = py_rand.randn(shape) torch_x = torch_matrix.float_tensor(py_x) py_y = py_func.tpow(py_x, power) torch_y = torch_func.tpow(torch_x, power) assert_close(py_y, torch_y, "tpow")
def test_reshape(): shape = (100,100) newshape = (5, 2000) py_rand.set_seed() py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) py_new = py_matrix.reshape(py_mat, newshape) torch_new = torch_matrix.reshape(torch_mat, newshape) assert_close(py_new, torch_new, "reshape")
def test_repeat(): shape = (100,) n_repeats = 5 py_rand.set_seed() py_x = py_rand.randn(shape) torch_x = torch_matrix.float_tensor(py_x) py_res = py_matrix.repeat(py_x, n_repeats) torch_res = torch_matrix.repeat(torch_x, n_repeats) assert_close(py_res, torch_res, "repeat")
def test_tround(): shape = (100,100) py_rand.set_seed() py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) py_round = py_matrix.tround(py_mat) torch_round = torch_matrix.tround(torch_mat) assert_close(py_round, torch_round, "tround")
def test_inv(): shape = (100, 100) py_rand.set_seed() py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) py_res = py_matrix.inv(py_mat) torch_res = torch_matrix.inv(torch_mat) # needs a lower tolerance to pass than other tests assert_close(py_res, torch_res, "inv", rtol=1e-4, atol=1e-5)
def test_batch_dot(): L = 10 N = 100 M = 50 v_shape = (L, N) W_shape = (N, M) h_shape = (L, M) py_rand.set_seed() py_v = py_rand.randn(v_shape) py_W = py_rand.randn(W_shape) py_h = py_rand.randn(h_shape) torch_v = torch_matrix.float_tensor(py_v) torch_W = torch_matrix.float_tensor(py_W) torch_h = torch_matrix.float_tensor(py_h) py_res = py_matrix.batch_dot(py_v, py_W, py_h) torch_res = torch_matrix.batch_dot(torch_v, torch_W, torch_h) assert_close(py_res, torch_res, "batch_dot")
def test_vstack(): # vector shape = (100,) py_rand.set_seed() py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) py_res = py_matrix.vstack([py_mat, py_mat]) torch_res = torch_matrix.vstack([torch_mat, torch_mat]) assert_close(py_res, torch_res, "vstack: vectors") # matrix shape = (100, 100) py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) py_res = py_matrix.vstack([py_mat, py_mat]) torch_res = torch_matrix.vstack([torch_mat, torch_mat]) assert_close(py_res, torch_res, "vstack: matrices")
def test_resample(): shape = (100, 50) n = 10 py_rand.set_seed() py_mat = py_rand.randn(shape) torch_mat = torch_matrix.float_tensor(py_mat) py_sample = py_matrix.resample(py_mat, n, replace = False) torch_sample = torch_matrix.resample(torch_mat, n, replace = False) # python backend and pytorch backend do not share a random number generator # therefore, we can only test that the shapes are the same assert py_matrix.shape(py_sample) == torch_matrix.shape(torch_sample), \ "python resample shape != torch resample shape"