現在，カーネル法による計算をjaxを用いて行うことを考えているのですが，題の通り，jaxとnumpyで計算結果が一致せず，その原因もわからないため，非常に困っています．

はじめに，カーネルとして，`K_train`が訓練データ同士のカーネル，`K_train_test`が訓練データとテストデータのカーネルとすると，推論時には訓練データのラベルを`y_train`，正則化のパラメータを`lamb`として以下のコードで求めます．

```

    pred = K_train_test.T @ np.linalg.inv(K_train + lamb * np.eye(K_train.shape[0])) @ y_train
    
```

以降，`np.linalg.inv(K_train + lamb + np.eye(K_train.shape[0]))`を`K_inv`とおきます．

  * まず`np.linalg.inv`では逆行列が求められるのに対して，`jax.numpy.linalg.inv`ではnanを返す点で計算結果が異なってしまいます．

  * 次に，`K_train_test.T @ K_inv`の計算結果もnumpyと異なっており，以下のような結果になっています．

```

    # numpy
    >>> K_train_test.T @ K_inv
    array([[-1.49965298e-03, -1.02000088e-03,  1.91625082e-04, ...,
            -3.85003951e-03, -7.90764352e-04,  3.35868994e-03],
           [-8.28401650e-04,  2.47865599e-03,  1.42909829e-03, ...,
             1.14423560e-02,  5.38865693e-04, -3.90940593e-03],
           [ 1.38003354e-03,  9.92750650e-03, -8.53253255e-03, ...,
             3.94552467e-03,  1.66892212e-02, -4.59781104e-03],
           ...,
           [ 1.69448515e-03, -1.69409449e-03, -8.55320507e-04, ...,
             1.26550822e-02,  6.71634316e-04,  1.34939273e-02],
           [ 8.37996740e-03,  7.80422727e-03, -5.00868738e-03, ...,
            -1.33403463e-03,  9.18144232e-05, -5.63108443e-02],
           [-9.93128179e-04, -1.23604874e-02,  8.20241083e-04, ...,
             1.85067963e-02,  1.67387109e-02,  7.49195804e-03]])
    
    # jax
    >>> K_train_test.T @ jax.device_put(K_inv)
    DeviceArray([[ 2.26987898e-03,  4.16899025e-02, -4.80616713e+00, ...,
                  -4.03137207e-02, -1.95407867e-03,  2.27355957e-03],
                 [ 7.90661573e-03, -4.59559262e-02, -4.97476482e+00, ...,
                  -2.04467773e-03, -1.02233887e-03, -3.71170044e-03],
                 [ 9.59877670e-03,  3.53920758e-02, -6.02732706e+00, ...,
                  -3.28140259e-02,  1.30906105e-02, -1.22079849e-02],
                 ...,
                 [ 9.36821103e-04,  6.76373541e-02, -4.46830130e+00, ...,
                  -1.94931030e-02,  5.78451157e-03,  1.83448792e-02],
                 [ 1.46173686e-02,  3.41360569e-02, -4.53566217e+00, ...,
                  -7.31468201e-03, -7.87258148e-04, -5.80210686e-02],
                 [ 7.50748813e-03, -4.72679436e-02, -5.52401400e+00, ...,
                   3.80363464e-02,  9.84239578e-03,  8.40377808e-03]],            dtype=float32)
    
```

ちなみに，分類問題を解いているのですが，numpyの計算結果ではAccuracyが0.97以上あるのに対して，jaxの方では0.1未満の結果となってしまい，jaxの計算結果がおかしいと考えています．

一方で，`K_inv @ y_train`の計算結果はnumpyでもjaxでも一致しているため，なおのこと原因がわからなくなっています．

以上のように，jaxではnumpyの結果と比べ結果が違うため，なにか見落としている部分がある気がするので，おかしな点があれば指摘していただけると助かります．
CC BY-SA 4.0
2022-11-26T05:04:01.213
92439
2022-11-26T05:43:27.167
2022-11-26T05:43:27.167
3060
55602
question
python numpy gpu
numpyによる計算とjaxによるもので計算結果が異なる