tf 2.0 force gradient

forcegrad.py

import tensorflow as tf
import numpy as np

@tf.custom_gradient
def force_grad_layer(x, grad_tensor):
  def grad(dy):
    return grad_tensor, tf.constant(0)
return tf.identity(x), grad

with tf.GradientTape() as tape:
  forced_grad = tf.constant(wanted_gradient_value)
  """variable description"""
  some_variable = tf.Variable(blahblah)
  """varibale description"""
  y = func_that_need_forced_gradient_value(last_layer)
  yy = force_grad_layer(y, forced_grad)
dyy_over_ds_with_dL_over_dyy_equals_wanted_gradient_value = tape.gradient(yy, some_variable)