Generating data with shape (50000, 1000)...
Data size: 400.0 MB
PyTorch CPU GMM
1.28 s ± 3.48 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
PyTorch GPU GMM
221 ms ± 1.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
NumPy GMM
1.6 s ± 5.42 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

x1 = tensor([[4.8761, 0.0000],
        [2.7456, 5.9371]], dtype=torch.float64)
x2 = tensor([[1., 0.],
        [0., 1.]], dtype=torch.float64), kwargs = {}
        x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
        # Torch tries to emulate NumPy 1 solve behavior by using batched 1-D solve
        # whenever
        # 1. x1.ndim - 1 == x2.ndim
        # 2. x1.shape[:-1] == x2.shape
        #
        # See linalg_solve_is_vector_rhs in
        # aten/src/ATen/native/LinearAlgebraUtils.h and
        # TORCH_META_FUNC(_linalg_solve_ex) in
        # aten/src/ATen/native/BatchLinearAlgebra.cpp in the PyTorch source code.
        #
        # The easiest way to work around this is to prepend a size 1 dimension to
        # x2, since x2 is already one dimension less than x1.
        #
        # See https://github.com/pytorch/pytorch/issues/52915
        if x2.ndim != 1 and x1.ndim - 1 == x2.ndim and x1.shape[:-1] == x2.shape:
            x2 = x2[None]
>       return torch.linalg.solve(x1, x2, **kwargs)
E       NameError: name 'torch' is not defined

kwargs     = {}
x1         = tensor([[4.8761, 0.0000],
        [2.7456, 5.9371]], dtype=torch.float64)
x2         = tensor([[1., 0.],
        [0., 1.]], dtype=torch.float64)