mwitiderrick/huber_loss.py

## huber_loss.py
logits = jnp.array([0.50,0.60,0.70,0.30,0.25])
labels = jnp.array([0.20,0.30,0.10,0.20,0.2])
optax.huber_loss(logits,labels)
# DeviceArray([0.045     , 0.045     , 0.17999998, 0.005     , 0.00125   ],            dtype=float32)
	logits = jnp.array([0.50,0.60,0.70,0.30,0.25])
	labels = jnp.array([0.20,0.30,0.10,0.20,0.2])
	optax.huber_loss(logits,labels)
	# DeviceArray([0.045 , 0.045 , 0.17999998, 0.005 , 0.00125 ], dtype=float32)