JaxPM/tests/helpers.py
Wassim Kabalan af29c4005d Format
2024-12-08 22:45:09 +01:00

13 lines
186 B
Python

import jax.numpy as jnp
def MSE(x, y):
return jnp.mean((x - y)**2)
def MSE_3D(x, y):
return ((x - y)**2).mean(axis=0)
def MSRE(x, y):
return jnp.mean(((x - y) / y)**2)