# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import numpy as np
import tensorflow as tf
from loguru import logger
from tensorflow.python.framework import ops
import numpy as np


class Dendrite:
    def __init__(self):
        pass
    
    def gradrpc(self, inputs, grad):
        logger.info('gradrpc {} {}', inputs, grad)
        return inputs

    def spike_grad(self, op, grad):
        logger.info('spike_grad')
        return tf.py_function(self.gradrpc, [op.inputs[0], grad], [tf.float32])

    def spikerpc(self, tensor):
        logger.info('spikerpc {}', tensor)
        return tensor

    def spike(self, inp, Tout, stateful=True, name=None, grad=None):
        rnd_name = 'PyFuncGrad' + str(np.random.randint(0, 1E+8))
        tf.RegisterGradient(rnd_name)(self.spike_grad) 
        g = tf.get_default_graph()
        with g.gradient_override_map({"PyFunc": rnd_name}):
            return tf.py_function(self.spikerpc, inp, Tout, name=name)

def pyfunc_test():
    
    g = tf.Graph()
    s = tf.Session()
    dendrite = Dendrite()

    x = tf.compat.v1.placeholder(tf.float32, shape=[1])
    
    w = tf.Variable(tf.ones_like(x))
    h = tf.multiply(x, w)
    
    y = dendrite.spike([h], [tf.float32])[0]
    
    d = tf.gradients(y, w)

    init = tf.compat.v1.global_variables_initializer()
    s.run(init)
    for out in s.run([y, d], feed_dict={x:[1]}):
        print (out)


#     # create data
#     x_data = tf.placeholder(dtype=tf.float32, shape=[None])
#     y_data = tf.placeholder(dtype=tf.float32, shape=[None])

#     w = tf.Variable(tf.constant([0.5]))
#     b = tf.Variable(tf.zeros([1]))

#     y1 = tf.multiply(w, x_data, name='y1')
#     y2 = spike([y1], [tf.float32], grad=spike_grad)[0]
#     y = tf.add(y2, b)

#     loss = tf.reduce_mean(tf.square(y - y_data))
#     optimizer = tf.train.GradientDescentOptimizer(0.01)
#     train = optimizer.minimize(loss)

#     print("Pyfunc grad", ops.get_gradient_function(y2.op))
#     with tf.Session() as sess:
#         sess.run(tf.global_variables_initializer())

#         for step in range(1):
#             #            ran = np.random.rand(115).astype(np.float32)
#             ran = np.ones((115)).astype(np.float32)
#             ans = ran * 1.5 + 3
#             dic = {x_data: ran, y_data: ans}
#             tt, yy, yy1= sess.run([train, y1, y2], feed_dict=dic)
#             if step % 1 == 0:
#                 print('step {}'.format(step))
#                 print('{}, {}'.format(w.eval(), b.eval()))

#         test = sess.run(y, feed_dict={x_data:[1]})
#         print('test = {}'.format(test))


if __name__ == '__main__':
    pyfunc_test()
# -



StarcoderdataJupyterScriptsDedupFiltered