উন্নত স্বয়ংক্রিয় পার্থক্য

TensorFlow.org এ দেখুন Google Colab-এ চালান GitHub-এ উৎস দেখুন নোটবুক ডাউনলোড করুন

গ্রেডিয়েন্টের ভূমিকা এবং স্বয়ংক্রিয় পার্থক্য নির্দেশিকাতে টেনসরফ্লোতে গ্রেডিয়েন্ট গণনা করার জন্য প্রয়োজনীয় সবকিছু অন্তর্ভুক্ত রয়েছে। এই নির্দেশিকাটি tf.GradientTape API-এর গভীর, কম সাধারণ বৈশিষ্ট্যগুলির উপর ফোকাস করে।

সেটআপ

import tensorflow as tf

import matplotlib as mpl
import matplotlib.pyplot as plt

mpl.rcParams['figure.figsize'] = (8, 6)

গ্রেডিয়েন্ট রেকর্ডিং নিয়ন্ত্রণ করা

স্বয়ংক্রিয় পার্থক্য নির্দেশিকাতে আপনি গ্রেডিয়েন্ট গণনা তৈরি করার সময় টেপ দ্বারা কোন ভেরিয়েবল এবং টেনসরগুলি দেখা হয় তা কীভাবে নিয়ন্ত্রণ করবেন তা দেখেছেন।

টেপে রেকর্ডিং ম্যানিপুলেট করার পদ্ধতিও রয়েছে।

রেকর্ডিং বন্ধ করুন

আপনি যদি গ্রেডিয়েন্ট রেকর্ডিং বন্ধ করতে চান, আপনি অস্থায়ীভাবে রেকর্ডিং স্থগিত করতে tf.GradientTape.stop_recording ব্যবহার করতে পারেন।

আপনি যদি আপনার মডেলের মাঝখানে একটি জটিল অপারেশনকে আলাদা করতে না চান তবে ওভারহেড কমাতে এটি কার্যকর হতে পারে। এটি একটি মেট্রিক বা একটি মধ্যবর্তী ফলাফল গণনা অন্তর্ভুক্ত করতে পারে:

x = tf.Variable(2.0)
y = tf.Variable(3.0)

with tf.GradientTape() as t:
  x_sq = x * x
  with t.stop_recording():
    y_sq = y * y
  z = x_sq + y_sq

grad = t.gradient(z, {'x': x, 'y': y})

print('dz/dx:', grad['x'])  # 2*x => 4
print('dz/dy:', grad['y'])

dz/dx: tf.Tensor(4.0, shape=(), dtype=float32)
dz/dy: None

স্ক্র্যাচ থেকে রেকর্ডিং রিসেট/শুরু করুন

আপনি যদি সম্পূর্ণভাবে শুরু করতে চান, tf.GradientTape.reset ব্যবহার করুন। কেবল গ্রেডিয়েন্ট টেপ ব্লক থেকে প্রস্থান করা এবং পুনরায় চালু করা সাধারণত পড়া সহজ, তবে টেপ ব্লক থেকে প্রস্থান করা কঠিন বা অসম্ভব হলে আপনি reset পদ্ধতি ব্যবহার করতে পারেন।

x = tf.Variable(2.0)
y = tf.Variable(3.0)
reset = True

with tf.GradientTape() as t:
  y_sq = y * y
  if reset:
    # Throw out all the tape recorded so far.
    t.reset()
  z = x * x + y_sq

grad = t.gradient(z, {'x': x, 'y': y})

print('dz/dx:', grad['x'])  # 2*x => 4
print('dz/dy:', grad['y'])

dz/dx: tf.Tensor(4.0, shape=(), dtype=float32)
dz/dy: None

নির্ভুলতার সাথে গ্রেডিয়েন্ট প্রবাহ বন্ধ করুন

উপরের গ্লোবাল টেপ কন্ট্রোলের বিপরীতে, tf.stop_gradient ফাংশন অনেক বেশি সুনির্দিষ্ট। এটি একটি নির্দিষ্ট পথ বরাবর প্রবাহিত থেকে গ্রেডিয়েন্টগুলি বন্ধ করতে ব্যবহার করা যেতে পারে, টেপটিতে অ্যাক্সেসের প্রয়োজন ছাড়াই:

x = tf.Variable(2.0)
y = tf.Variable(3.0)

with tf.GradientTape() as t:
  y_sq = y**2
  z = x**2 + tf.stop_gradient(y_sq)

grad = t.gradient(z, {'x': x, 'y': y})

print('dz/dx:', grad['x'])  # 2*x => 4
print('dz/dy:', grad['y'])

dz/dx: tf.Tensor(4.0, shape=(), dtype=float32)
dz/dy: None

কাস্টম গ্রেডিয়েন্ট

কিছু ক্ষেত্রে, আপনি ডিফল্ট ব্যবহার করার পরিবর্তে গ্রেডিয়েন্টগুলি ঠিক কীভাবে গণনা করা হয় তা নিয়ন্ত্রণ করতে চাইতে পারেন। এই পরিস্থিতিতে অন্তর্ভুক্ত:

  1. আপনি লিখছেন একটি নতুন অপের জন্য কোন সংজ্ঞায়িত গ্রেডিয়েন্ট নেই.
  2. ডিফল্ট গণনা সংখ্যাগতভাবে অস্থির।
  3. আপনি ফরোয়ার্ড পাস থেকে একটি ব্যয়বহুল গণনা ক্যাশে করতে চান।
  4. আপনি গ্রেডিয়েন্ট পরিবর্তন না করে একটি মান পরিবর্তন করতে চান (উদাহরণস্বরূপ, tf.clip_by_value বা tf.math.round ব্যবহার করে)।

প্রথম ক্ষেত্রে, একটি নতুন বিকল্প লিখতে আপনি আপনার নিজস্ব সেট আপ করতে tf.RegisterGradient ব্যবহার করতে পারেন (বিশদ বিবরণের জন্য API ডক্স দেখুন)। (দ্রষ্টব্য যে গ্রেডিয়েন্ট রেজিস্ট্রি বিশ্বব্যাপী, তাই সাবধানতার সাথে এটি পরিবর্তন করুন।)

পরবর্তী তিনটি ক্ষেত্রে, আপনি tf.custom_gradient ব্যবহার করতে পারেন।

এখানে একটি উদাহরণ রয়েছে যা মধ্যবর্তী গ্রেডিয়েন্টে tf.clip_by_norm প্রয়োগ করে:

# Establish an identity operation, but clip during the gradient pass.
@tf.custom_gradient
def clip_gradients(y):
  def backward(dy):
    return tf.clip_by_norm(dy, 0.5)
  return y, backward

v = tf.Variable(2.0)
with tf.GradientTape() as t:
  output = clip_gradients(v * v)
print(t.gradient(output, v))  # calls "backward", which clips 4 to 2

tf.Tensor(2.0, shape=(), dtype=float32)

আরো বিস্তারিত জানার জন্য tf.custom_gradient decorator API ডক্স দেখুন।

SavedModel-এ কাস্টম গ্রেডিয়েন্ট

কাস্টম গ্রেডিয়েন্টগুলি tf.saved_model.SaveOptions(experimental_custom_gradients=True) বিকল্পটি ব্যবহার করে SavedModel-এ সংরক্ষণ করা যেতে পারে।

SavedModel-এ সেভ করার জন্য, গ্রেডিয়েন্ট ফাংশন অবশ্যই ট্রেসযোগ্য হতে হবে (আরো জানতে, tf.function গাইডের সাথে আরও ভাল পারফরম্যান্স দেখুন)।

class MyModule(tf.Module):

  @tf.function(input_signature=[tf.TensorSpec(None)])
  def call_custom_grad(self, x):
    return clip_gradients(x)

model = MyModule()

tf.saved_model.save(
    model,
    'saved_model',
    options=tf.saved_model.SaveOptions(experimental_custom_gradients=True))

# The loaded gradients will be the same as the above example.
v = tf.Variable(2.0)
loaded = tf.saved_model.load('saved_model')
with tf.GradientTape() as t:
  output = loaded.call_custom_grad(v * v)
print(t.gradient(output, v))

INFO:tensorflow:Assets written to: saved_model/assets
tf.Tensor(2.0, shape=(), dtype=float32)

উপরের উদাহরণ সম্পর্কে একটি নোট: আপনি যদি উপরের tf.saved_model.SaveOptions(experimental_custom_gradients=False) দিয়ে প্রতিস্থাপন করার চেষ্টা করেন, তাহলে গ্রেডিয়েন্ট লোড করার সময়ও একই ফলাফল দেবে। কারণ হল যে গ্রেডিয়েন্ট রেজিস্ট্রিতে এখনও ফাংশন call_custom_op এ ব্যবহৃত কাস্টম গ্রেডিয়েন্ট রয়েছে। যাইহোক, যদি আপনি কাস্টম গ্রেডিয়েন্ট ছাড়া সংরক্ষণ করার পরে রানটাইম পুনরায় চালু করেন, তাহলে tf.GradientTape এর অধীনে লোড করা মডেলটি চালালে ত্রুটি হবে: LookupError: No gradient defined for operation 'IdentityN' (op type: IdentityN) ।

একাধিক টেপ

একাধিক টেপ নির্বিঘ্নে যোগাযোগ করে।

উদাহরণস্বরূপ, এখানে প্রতিটি টেপ বিভিন্ন টেনসরের সেট দেখে:

x0 = tf.constant(0.0)
x1 = tf.constant(0.0)

with tf.GradientTape() as tape0, tf.GradientTape() as tape1:
  tape0.watch(x0)
  tape1.watch(x1)

  y0 = tf.math.sin(x0)
  y1 = tf.nn.sigmoid(x1)

  y = y0 + y1

  ys = tf.reduce_sum(y)

tape0.gradient(ys, x0).numpy()   # cos(x) => 1.0

1.0

tape1.gradient(ys, x1).numpy()   # sigmoid(x1)*(1-sigmoid(x1)) => 0.25

0.25

উচ্চ ক্রম গ্রেডিয়েন্ট

tf.GradientTape কনটেক্সট ম্যানেজারের ভিতরের কাজগুলি স্বয়ংক্রিয় পার্থক্যের জন্য রেকর্ড করা হয়। যদি গ্রেডিয়েন্টগুলি সেই প্রসঙ্গে গণনা করা হয়, তবে গ্রেডিয়েন্ট গণনাটিও রেকর্ড করা হয়। ফলস্বরূপ, সঠিক একই API উচ্চ-ক্রম গ্রেডিয়েন্টের জন্যও কাজ করে।

উদাহরণ স্বরূপ:

x = tf.Variable(1.0)  # Create a Tensorflow variable initialized to 1.0

with tf.GradientTape() as t2:
  with tf.GradientTape() as t1:
    y = x * x * x

  # Compute the gradient inside the outer `t2` context manager
  # which means the gradient computation is differentiable as well.
  dy_dx = t1.gradient(y, x)
d2y_dx2 = t2.gradient(dy_dx, x)

print('dy_dx:', dy_dx.numpy())  # 3 * x**2 => 3.0
print('d2y_dx2:', d2y_dx2.numpy())  # 6 * x => 6.0

dy_dx: 3.0
d2y_dx2: 6.0

যদিও এটি আপনাকে একটি স্কেলার ফাংশনের দ্বিতীয় ডেরিভেটিভ দেয়, এই প্যাটার্নটি একটি হেসিয়ান ম্যাট্রিক্স তৈরি করতে সাধারণীকরণ করে না, যেহেতু tf.GradientTape.gradient শুধুমাত্র একটি স্কেলারের গ্রেডিয়েন্ট গণনা করে। একটি হেসিয়ান ম্যাট্রিক্স তৈরি করতে, জ্যাকোবিয়ান বিভাগের অধীনে হেসিয়ান উদাহরণে যান।

" tf.GradientTape.gradient এ নেস্টেড কল" হল একটি ভাল প্যাটার্ন যখন আপনি একটি গ্রেডিয়েন্ট থেকে একটি স্কেলার গণনা করছেন এবং তারপরে প্রাপ্ত স্কেলারটি একটি দ্বিতীয় গ্রেডিয়েন্ট গণনার উত্স হিসাবে কাজ করে, যেমনটি নিম্নলিখিত উদাহরণে।

উদাহরণ: ইনপুট গ্রেডিয়েন্ট নিয়মিতকরণ

অনেক মডেল "বিরোধী উদাহরণ" এর জন্য সংবেদনশীল। কৌশলগুলির এই সংগ্রহটি মডেলের আউটপুটকে বিভ্রান্ত করতে মডেলের ইনপুট পরিবর্তন করে। সহজতম বাস্তবায়ন—যেমন ফাস্ট গ্রেডিয়েন্ট সাইনড মেথড অ্যাটাক ব্যবহার করে অ্যাডভারসারিয়াল উদাহরণ —ইনপুটের ক্ষেত্রে আউটপুটের গ্রেডিয়েন্টের সাথে একক পদক্ষেপ নেয়; "ইনপুট গ্রেডিয়েন্ট"।

প্রতিকূল উদাহরণগুলির দৃঢ়তা বাড়ানোর একটি কৌশল হল ইনপুট গ্রেডিয়েন্ট নিয়মিতকরণ (ফিনলে এবং ওবারম্যান, 2019), যা ইনপুট গ্রেডিয়েন্টের মাত্রা কমানোর চেষ্টা করে। যদি ইনপুট গ্রেডিয়েন্ট ছোট হয়, তাহলে আউটপুটে পরিবর্তনটিও ছোট হওয়া উচিত।

নীচে ইনপুট গ্রেডিয়েন্ট রেগুলারাইজেশনের একটি নির্বোধ বাস্তবায়ন। বাস্তবায়ন হল:

  1. একটি অভ্যন্তরীণ টেপ ব্যবহার করে ইনপুটের সাপেক্ষে আউটপুটের গ্রেডিয়েন্ট গণনা করুন।
  2. সেই ইনপুট গ্রেডিয়েন্টের মাত্রা গণনা করুন।
  3. মডেলের সাপেক্ষে সেই মাত্রার গ্রেডিয়েন্ট গণনা করুন।
x = tf.random.normal([7, 5])

layer = tf.keras.layers.Dense(10, activation=tf.nn.relu)

with tf.GradientTape() as t2:
  # The inner tape only takes the gradient with respect to the input,
  # not the variables.
  with tf.GradientTape(watch_accessed_variables=False) as t1:
    t1.watch(x)
    y = layer(x)
    out = tf.reduce_sum(layer(x)**2)
  # 1. Calculate the input gradient.
  g1 = t1.gradient(out, x)
  # 2. Calculate the magnitude of the input gradient.
  g1_mag = tf.norm(g1)

# 3. Calculate the gradient of the magnitude with respect to the model.
dg1_mag = t2.gradient(g1_mag, layer.trainable_variables)

[var.shape for var in dg1_mag]

[TensorShape([5, 10]), TensorShape([10])]

জ্যাকোবিয়ানরা

পূর্ববর্তী সমস্ত উদাহরণগুলি কিছু উত্স টেনসর(গুলি) এর সাপেক্ষে একটি স্কেলার লক্ষ্যের গ্রেডিয়েন্ট গ্রহণ করেছে।

জ্যাকোবিয়ান ম্যাট্রিক্স একটি ভেক্টর ভ্যালুড ফাংশনের গ্রেডিয়েন্টকে উপস্থাপন করে। প্রতিটি সারিতে ভেক্টরের উপাদানগুলির একটির গ্রেডিয়েন্ট রয়েছে।

tf.GradientTape.jacobian পদ্ধতি আপনাকে একটি জ্যাকোবিয়ান ম্যাট্রিক্স দক্ষতার সাথে গণনা করতে দেয়।

মনে রাখবেন যে:

  • gradient মতো: sources যুক্তি একটি টেনসর বা টেনসরের একটি ধারক হতে পারে।
  • gradient বিপরীতে: target টেনসর অবশ্যই একটি একক টেনসর হতে হবে।

স্কেলার উত্স

প্রথম উদাহরণ হিসাবে, এখানে একটি স্কেলার-উৎস সম্পর্কিত ভেক্টর-টার্গেটের জ্যাকোবিয়ান।

x = tf.linspace(-10.0, 10.0, 200+1)
delta = tf.Variable(0.0)

with tf.GradientTape() as tape:
  y = tf.nn.sigmoid(x+delta)

dy_dx = tape.jacobian(y, delta)

আপনি যখন একটি স্কেলারের সাথে জ্যাকোবিয়ান নেন তখন ফলাফলটি লক্ষ্যের আকৃতি ধারণ করে এবং উৎসের সাপেক্ষে প্রতিটি উপাদানের গ্রেডিয়েন্ট দেয়:

print(y.shape)
print(dy_dx.shape)

(201,)
(201,)

plt.plot(x.numpy(), y, label='y')
plt.plot(x.numpy(), dy_dx, label='dy/dx')
plt.legend()
_ = plt.xlabel('x')

png

টেনসর উৎস

ইনপুটটি স্কেলার বা টেনসর হোক না কেন, tf.GradientTape.jacobian দক্ষতার সাথে লক্ষ্য(গুলি) এর প্রতিটি উপাদানের সাপেক্ষে উত্সের প্রতিটি উপাদানের গ্রেডিয়েন্ট গণনা করে।

উদাহরণস্বরূপ, এই স্তরটির আউটপুটের একটি আকৃতি রয়েছে (10, 7) :

x = tf.random.normal([7, 5])
layer = tf.keras.layers.Dense(10, activation=tf.nn.relu)

with tf.GradientTape(persistent=True) as tape:
  y = layer(x)

y.shape

TensorShape([7, 10])

এবং স্তরটির কার্নেলের আকৃতি হল (5, 10) :

layer.kernel.shape

TensorShape([5, 10])

কার্নেলের সাপেক্ষে আউটপুটের জ্যাকোবিয়ানের আকৃতি হল সেই দুটি আকৃতি একত্রে সংযুক্ত:

j = tape.jacobian(y, layer.kernel)
j.shape

TensorShape([7, 10, 5, 10])

যদি আপনি লক্ষ্যের মাত্রার উপর যোগ করেন, তাহলে আপনার যোগফলের গ্রেডিয়েন্ট বাকি থাকবে যা tf.GradientTape.gradient দ্বারা গণনা করা হবে :

g = tape.gradient(y, layer.kernel)
print('g.shape:', g.shape)

j_sum = tf.reduce_sum(j, axis=[0, 1])
delta = tf.reduce_max(abs(g - j_sum)).numpy()
assert delta < 1e-3
print('delta:', delta)

g.shape: (5, 10)
delta: 2.3841858e-07

উদাহরণ: হেসিয়ান

যদিও tf.GradientTape একটি হেসিয়ান ম্যাট্রিক্স নির্মাণের জন্য একটি সুস্পষ্ট পদ্ধতি দেয় না, tf.GradientTape.jacobian পদ্ধতি ব্যবহার করে একটি তৈরি করা সম্ভব।

x = tf.random.normal([7, 5])
layer1 = tf.keras.layers.Dense(8, activation=tf.nn.relu)
layer2 = tf.keras.layers.Dense(6, activation=tf.nn.relu)

with tf.GradientTape() as t2:
  with tf.GradientTape() as t1:
    x = layer1(x)
    x = layer2(x)
    loss = tf.reduce_mean(x**2)

  g = t1.gradient(loss, layer1.kernel)

h = t2.jacobian(g, layer1.kernel)

print(f'layer.kernel.shape: {layer1.kernel.shape}')
print(f'h.shape: {h.shape}')

layer.kernel.shape: (5, 8)
h.shape: (5, 8, 5, 8)

একটি নিউটনের পদ্ধতি পদক্ষেপের জন্য এই হেসিয়ান ব্যবহার করতে, আপনি প্রথমে এটির অক্ষগুলিকে একটি ম্যাট্রিক্সে সমতল করবেন এবং গ্রেডিয়েন্টটিকে একটি ভেক্টরে সমতল করবেন:

n_params = tf.reduce_prod(layer1.kernel.shape)

g_vec = tf.reshape(g, [n_params, 1])
h_mat = tf.reshape(h, [n_params, n_params])

হেসিয়ান ম্যাট্রিক্স প্রতিসম হওয়া উচিত:

def imshow_zero_center(image, **kwargs):
  lim = tf.reduce_max(abs(image))
  plt.imshow(image, vmin=-lim, vmax=lim, cmap='seismic', **kwargs)
  plt.colorbar()

imshow_zero_center(h_mat)

png

নিউটনের পদ্ধতি আপডেট ধাপ নিচে দেখানো হয়েছে:

eps = 1e-3
eye_eps = tf.eye(h_mat.shape[0])*eps

# X(k+1) = X(k) - (∇²f(X(k)))^-1 @ ∇f(X(k))
# h_mat = ∇²f(X(k))
# g_vec = ∇f(X(k))
update = tf.linalg.solve(h_mat + eye_eps, g_vec)

# Reshape the update and apply it to the variable.
_ = layer1.kernel.assign_sub(tf.reshape(update, layer1.kernel.shape))

যদিও এটি একটি একক tf.Variable এর জন্য তুলনামূলকভাবে সহজ, একটি নন-তুচ্ছ মডেলে এটি প্রয়োগ করার জন্য একাধিক ভেরিয়েবল জুড়ে একটি সম্পূর্ণ হেসিয়ান তৈরি করতে সতর্ক সংমিশ্রণ এবং স্লাইসিং প্রয়োজন হবে।

ব্যাচ জ্যাকোবিয়ান

কিছু ক্ষেত্রে, আপনি উৎসের স্ট্যাকের সাপেক্ষে প্রতিটি লক্ষ্যের স্ট্যাকের জ্যাকোবিয়ান নিতে চান, যেখানে প্রতিটি লক্ষ্য-উৎস জোড়ার জন্য জ্যাকোবিয়ানরা স্বাধীন।

উদাহরণস্বরূপ, এখানে ইনপুট x আকৃতির (batch, ins) এবং আউটপুট y আকৃতির (batch, outs) :

x = tf.random.normal([7, 5])

layer1 = tf.keras.layers.Dense(8, activation=tf.nn.elu)
layer2 = tf.keras.layers.Dense(6, activation=tf.nn.elu)

with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:
  tape.watch(x)
  y = layer1(x)
  y = layer2(y)

y.shape

TensorShape([7, 6])

x এর সাপেক্ষে y এর পূর্ণ Jacobian-এর একটি আকৃতি আছে (batch, ins, batch, outs) , এমনকি যদি আপনি শুধুমাত্র চান (batch, ins, outs) :

j = tape.jacobian(y, x)
j.shape

TensorShape([7, 6, 7, 5])

যদি স্ট্যাকের প্রতিটি আইটেমের গ্রেডিয়েন্ট স্বাধীন হয়, তাহলে এই টেনসরের প্রতিটি (batch, batch) স্লাইস একটি তির্যক ম্যাট্রিক্স:

imshow_zero_center(j[:, 0, :, 0])
_ = plt.title('A (batch, batch) slice')

png

def plot_as_patches(j):
  # Reorder axes so the diagonals will each form a contiguous patch.
  j = tf.transpose(j, [1, 0, 3, 2])
  # Pad in between each patch.
  lim = tf.reduce_max(abs(j))
  j = tf.pad(j, [[0, 0], [1, 1], [0, 0], [1, 1]],
             constant_values=-lim)
  # Reshape to form a single image.
  s = j.shape
  j = tf.reshape(j, [s[0]*s[1], s[2]*s[3]])
  imshow_zero_center(j, extent=[-0.5, s[2]-0.5, s[0]-0.5, -0.5])

plot_as_patches(j)
_ = plt.title('All (batch, batch) slices are diagonal')

png

পছন্দসই ফলাফল পেতে, আপনি ডুপ্লিকেট batch মাত্রা যোগ করতে পারেন, অথবা tf.einsum ব্যবহার করে কর্ণ নির্বাচন করতে পারেন:

j_sum = tf.reduce_sum(j, axis=2)
print(j_sum.shape)
j_select = tf.einsum('bxby->bxy', j)
print(j_select.shape)

(7, 6, 5)
(7, 6, 5)

প্রথম স্থানে অতিরিক্ত মাত্রা ছাড়াই গণনা করা অনেক বেশি কার্যকর হবে। tf.GradientTape.batch_jacobian পদ্ধতি ঠিক এটি করে:

jb = tape.batch_jacobian(y, x)
jb.shape

WARNING:tensorflow:5 out of the last 5 calls to <function pfor.<locals>.f at 0x7f7d601250e0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
TensorShape([7, 6, 5])

error = tf.reduce_max(abs(jb - j_sum))
assert error < 1e-3
print(error.numpy())

0.0

x = tf.random.normal([7, 5])

layer1 = tf.keras.layers.Dense(8, activation=tf.nn.elu)
bn = tf.keras.layers.BatchNormalization()
layer2 = tf.keras.layers.Dense(6, activation=tf.nn.elu)

with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:
  tape.watch(x)
  y = layer1(x)
  y = bn(y, training=True)
  y = layer2(y)

j = tape.jacobian(y, x)
print(f'j.shape: {j.shape}')

WARNING:tensorflow:6 out of the last 6 calls to <function pfor.<locals>.f at 0x7f7cf062fa70> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
j.shape: (7, 6, 7, 5)

plot_as_patches(j)

_ = plt.title('These slices are not diagonal')
_ = plt.xlabel("Don't use `batch_jacobian`")

png

এই ক্ষেত্রে, batch_jacobian এখনও রান করে এবং প্রত্যাশিত আকারের সাথে কিছু ফেরত দেয়, কিন্তু এর বিষয়বস্তুর একটি অস্পষ্ট অর্থ রয়েছে:

jb = tape.batch_jacobian(y, x)
print(f'jb.shape: {jb.shape}')

jb.shape: (7, 6, 5)