fix loss functions after math module reorganization
@@ -9,10 +9,10 @@ from phi import struct
from . import _extrapolation as extrapolation
from ._extrapolation import Extrapolation
from . import _functions as math
-from ._shape import CHANNEL_DIM, spatial_shape, channel_shape, Shape
+from ._shape import CHANNEL_DIM, spatial_shape, channel_shape, Shape, shape
from ._functions import broadcast_op, batch_stack, channel_stack
from ._tensors import tensor, Tensor, TensorStack, NativeTensor
-from ._tensors import Tensor
+from ._tensors import Tensor, shapeof
from ._config import GLOBAL_AXIS_ORDER
from .backend.tensorop import collapsed_gather_nd
@@ -70,8 +70,8 @@ def l1_loss(tensor: Tensor, batch_norm=True, reduce_batches=True):
else:
total_loss = math.sum_(math.abs(tensor), axis=list(range(1, len(tensor.shape))))
if batch_norm and reduce_batches:
- batch_size = math.shape(tensor)[0]
- return math.div(total_loss, math.to_float(batch_size))
+ batch_size = shapeof(tensor)[0]
+ return math.divide_no_nan(total_loss, math.to_float(batch_size))
else:
return total_loss
@@ -88,8 +88,8 @@ def l_n_loss(tensor: Tensor, n: int, batch_norm=True):
return sum(l_n_loss(tensor, n, batch_norm) for tensor in all_tensors)
total_loss = math.sum_(tensor ** n) / n
if batch_norm:
- batch_size = math.shape(tensor)[0]
- return math.div(total_loss, math.to_float(batch_size))
+ batch_size = shapeof(tensor)[0]
+ return math.divide_no_nan(total_loss, math.to_float(batch_size))
else:
return total_loss
