import tensorflow as tf

tf_gradient_function = tf_gradients_lib.gradients

# ISSUE: https://github.com/cybertronai/gradient-checkpointing/issues/38
def tf_gradients(ys, *args, **kwargs):
    """Decorate tf.gradients calls with explicit device placement to avoid memory
    leaks when splitting model across multiple GPUs"""
    source = ys[0] if isinstance(ys, (list, tuple)) else ys
    device = source.op.node_def.device if isinstance(source, tf.Tensor) else None
    with tf.device(device):
        return tf_gradient_function(ys, *args, **kwargs)


MIN_CHECKPOINT_NODE_SIZE=1024    # use lower value during testing

# specific versions we can use to do process-wide replacement of tf.gradients
def gradients_speed(ys, xs, grad_ys=None, **kwargs):
