参考

TensorFlow公式ガイド Recurrent Neural Networks (RNN) with Keras
tf.keras.layers.RNN のレファレンス
 TensorFlow公式ガイド Advanced Automatic differentiation

準備

import numpy as np
import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers

tf.keras.backend.set_floatx('float64')
mydtype = tf.float64

倍精度で計算します．

RNN cell としてダイナミクスを定義

カスタマイズしたRNNモデルを作るには，まず1ステップ分を計算するためのクラス（RNNセルと呼ばれる）を定義しインスタンスを作ります．
今回このセルを力学系の1ステップ分を計算するものとして扱います．

cellはstate_size, output_sizeの属性とcallメソッドを持っておく必要があります．今回はRNNの内部状態stateをそのまま力学系の状態変数として使うので，state_sizeは力学系の次元の値に設定します．出力は，内部状態の値そのままを出力するので，output_sizeも同様に力学系の次元の値にします．

callメソッドで1ステップ分の値を記述します．外部入力inputsは引数として必要ですが，今回は使いません．
レファレンスには'A call(input_at_t, states_at_t) method, returning (output_at_t, states_at_t_plus_1). 'と書いてあるとおり，ある時刻での入力と内部状態を引数で受け取り，外部出力と次の時刻の内部状態を返すメソッドとして作成します．それぞれはバッチ処理を前提にしているので，二次元配列となり，1次元めはバッチ内のサンプル数になります．

callの中身では，エノン写像$f(x,y) = (1-ax^2+y, bx) $ を計算しています．バッチ処理を前提にしています．

class HenonCell(keras.layers.Layer):
  '''
  This class calculates one step of Henon-map,
  used as a RNN cell by Keras RNN layer.
  '''
  def __init__(self, a=1.2, b=0.4, **kwargs):

    # parameter
    self.a = tf.constant(a, dtype=mydtype)
    self.b = tf.constant(b, dtype=mydtype)
    # RNN custom cell must have state_size and output_size
    self.state_size = 2 # 2-dimension
    self.output_size = 2
      
    super(HenonCell, self).__init__(**kwargs)

  def build(self, input_shape):
    self.batch_size = input_shape[0]

  def call(self, inputs, states):
    ''' states is a tuple, which only has a (batch_size, 2)-array.
        inputs are ignored.
    '''
    # print(type(states)) # tuple
    # print(states[0].shape) # (2,)
    x =states[0]
    x_new = 1.0 - self.a *x[:,0]*x[:,0] + x[:,1]
    y_new = self.b * x[:,0]
    X_new = tf.stack([x_new,y_new],axis=1)
    return X_new,[X_new] # return output and states.

1ステップの計算

cellのインスタンスに値を入れれば，写像を計算することができます．hcの返り値は出力と内部状態のタプルです．出力をx_newに入れています．

if __name__ == '__main__':

  t_length = 10000 #軌道の長さ
  batch_size = 1

  hc = HenonCell(a=1.4, b=0.3)# cell
  x0 = tf.constant([[0.1, 0.1]], dtype=mydtype)  # shape is (1,2)
  x_new, _ = hc(None, [x0])  # one-step
  print(x_new)

RNNレイヤーのクラスを作る

cellインスタンスを利用してRNNレイヤーのクラスを作ります．これが，軌道を一度に計算するために使えます．

  henon_rnn = layers.RNN(hc, return_sequences=True)

軌道の計算

軌道を計算します．ダミーでinputをいれています．この配列の1次元めのサイズ（時間方向の長さ）で，何ステップ計算するかが決まります．
初期値は RNNをcallするときにinitial_stateで渡すことができます．


 #making dummy input, which determines sequence length.
  dummy_input = tf.zeros(shape=[1, t_length, 2], dtype=mydtype)

  #initial state
  x0 = tf.constant(0.1*np.random.randn(1, 2), dtype=mydtype)
  # run
  X = henon_rnn(dummy_input, initial_state=x0)

  print(f'output sequence shape:{X.shape}')# (1,10000,2)

結果の確認

  import matplotlib.pyplot as plt

  plt.plot(X[0, :, 0], X[0, :, 1], '.', markersize=0.1)

Jacobian

tf.GradientTapeを使えば，勾配を計算する場合と同じように，Jacobianを計算することができます．軌道の各点でのJacobianを一度に計算するには batch_jacobian を使います．(公式ガイド参照)


  X = henon_rnn(dummy_input, initial_state=x0)
  
  #reshape (bs, t_length,2) -> (bs*t_length,2)
  X = tf.reshape(X, [-1, 2])  # (batch_size*t_length, 2)

  with tf.GradientTape() as tape:
    tape.watch(X)
    X_next, _ = hc(None, [X])
  X_next.shape

  Js = tape.batch_jacobian(X_next, X)
  print(Js.shape)  # (1000,2,2)
  print(Js.numpy()[0,:,:])

これはできない

RNNを動かして出てくる軌道Xから直接Jacobianを計算することはできません．

  with tf.GradientTape(persistent=True) as tape:
    tape.watch(x0)
    X = henon_rnn(dummy_input, initial_state=x0)

  Jx = tape.batch_jacobian(X[:,1,:], X[:,0,:]) # 答えはNoneになってしまう．