In [1]:

    

import tensorflow as tf
import numpy as np
import pylab as plt

@tf.RegisterGradient('WrapGrad')
def _wrap_grad(op,grad):
    phi = op.inputs[0]
    return tf.ones_like(phi)*grad

def wrap(phi):
    out = tf.atan2(tf.sin(phi),tf.cos(phi))
    with tf.get_default_graph().gradient_override_map({'Identity': 'WrapGrad'}):
        return tf.identity(out)


    

    




/home/josh/anaconda3/envs/kerastf/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.
  from ._conv import register_converters as _register_converters