<reponame>VulcanClimateModeling/fv3net
import tensorflow as tf
from typing import Callable, Mapping, Tuple

TensorDict = Mapping[str, tf.Tensor]
LossFunction = Callable[[TensorDict, TensorDict], Tuple[tf.Tensor, TensorDict]]

import tensorflow as tf
from typing import Callable, Mapping, Tuple

TensorDict = Mapping[str, tf.Tensor]
LossFunction = Callable[[TensorDict, TensorDict], Tuple[tf.Tensor, TensorDict]]