# 機能設計書 38-前方自動微分（Forward-mode AD）

## 概要

本ドキュメントは、TensorFlowにおける前方自動微分（Forward-mode Automatic Differentiation）機能の設計を記載する。`tf.autodiff.ForwardAccumulator` クラスを中心とし、ヤコビアン-ベクトル積（JVP）を計算する機構であり、`tensorflow/python/eager/forwardprop.py` に実装される。

### 本機能の処理概要

**業務上の目的・背景**：逆方向自動微分（GradientTape）はベクトル-ヤコビアン積（VJP）を計算するのに対し、前方自動微分はヤコビアン-ベクトル積（JVP）を計算する。出力が多く入力が少ない関数の場合、前方モードが効率的である。また、中間活性を保持しないため、メモリ効率が優れている。

**機能の利用シーン**：
- ヘシアン-ベクトル積の計算（forward-over-backward）
- 出力次元が入力次元より大きい関数の微分
- メモリ効率が重要な場合

**主要な処理内容**：
1. `ForwardAccumulator` コンテキストマネージャでprimal-tangentペアを管理
2. コンテキスト内で実行されるOpに対して自動的にJVPを計算
3. `jvp()` メソッドで任意のprimalテンソルに対応するtangentを取得
4. 特殊ケース（Identity, ReadVariableOp）は専用のJVP関数で処理
5. 一般ケースは逆方向勾配関数の転置により前方微分を計算

**関連システム・外部連携**：C++ `pywrap_tfe` モジュールのJVPディスパッチ機構、`tracing_compilation` による関数トレース、GradientTape（転置計算用）。

**権限による制御**：特になし。

## 関連画面

| 画面No | 画面名 | 関連種別 | 関連する操作・処理 |
|--------|--------|----------|------------------|
| - | 該当なし | - | 画面機能マッピングに関連画面なし |

## 機能種別

計算処理（前方自動微分）

## 入力仕様

### ForwardAccumulatorコンストラクタ

| パラメータ名 | 型 | 必須 | 説明 | バリデーション |
|-------------|-----|-----|------|---------------|
| primals | Tensor / Variable / list | Yes | 監視対象のテンソル | 同一テンソルの複数指定は不可 |
| tangents | Tensor / list | Yes | primalsに対応するJVPベクトル | primalsと同じ構造・形状 |

### jvp() メソッド

| パラメータ名 | 型 | 必須 | 説明 | バリデーション |
|-------------|-----|-----|------|---------------|
| primals | Tensor | Yes | JVPを取得したいprimalテンソル | ForwardAccumulator内で計算されたテンソル |

## 出力仕様

### 出力データ

| 項目名 | 型 | 説明 |
|--------|-----|------|
| tangent | Tensor / None | primalに対応するJVP（接線ベクトル）。対応がない場合None |

### 出力先

呼び出し元にtangentテンソルとして返却。

## 処理フロー

### 処理シーケンス

```
1. ForwardAccumulator.__init__: primals-tangentsペアを保存
2. ForwardAccumulator.__enter__: アキュムレータをスタックにプッシュ
3. Op実行時: _jvp_dispatch()が呼ばれる
   └─ _TRACE_COUNTが_TRACE_COUNT_LIMIT(32)未満: exact shape tracing
   └─ _TRACE_COUNT_LIMIT以上: relaxed shape tracing
   └─ _jvp_helper_wrapper -> _jvp_helper
4. _jvp_helper:
   └─ 特殊ケース（Identity, ReadVariableOp）: 直接tangent返却
   └─ 一般ケース: GradientTapeで逆方向勾配を記録 -> transposeでJVP計算
5. ForwardAccumulator.__exit__: アキュムレータをポップ
6. jvp(): primalに対応するtangentを取得
```

### フローチャート

```mermaid
flowchart TD
    A[ForwardAccumulator.__enter__] --> B[primals-tangentsペア登録]
    B --> C[コンテキスト内のOp実行]
    C --> D[_jvp_dispatch呼び出し]
    D --> E{特殊ケース?}
    E -->|Identity/ReadVariableOp| F[直接tangent返却]
    E -->|一般| G[_jvp_helper]
    G --> H[GradientTapeで逆方向勾配記録]
    H --> I[transpose_tapeでJVP計算]
    F --> J[tangent保存]
    I --> J
    J --> K[jvp()でtangent取得]
```

## ビジネスルール

### 業務ルール

| ルールNo | ルール名 | 内容 | 適用条件 |
|---------|---------|------|---------|
| BR-38-1 | primal一意性 | 同一テンソルをprimalsに複数回指定すると例外 | __init__時 |
| BR-38-2 | ネスト順序 | ForwardAccumulator同士のネストは作成順に適用。内側は外側のJVPを見ない | ネスト時 |
| BR-38-3 | GradientTapeとのネスト | より深くネストされたGradientTapeは外側のForwardAccumulatorのJVPを無視 | forward-over-backward計算時 |
| BR-38-4 | トレースカウント制限 | _TRACE_COUNT_LIMIT=32を超えると形状緩和トレーシングに切り替え | _jvp_dispatch時 |
| BR-38-5 | 特殊ケースJVP | Identity, ReadVariableOpはarray_ops.identityで直接tangentを返す | 該当Op実行時 |

### 計算ロジック

一般的なJVP計算（逆方向勾配の転置）:
```
1. 内側のGradientTapeで逆方向勾配関数を記録:
   backfunc_tape.watch(trainable_inputs)
   grads = backfunc_tape.gradient(outputs, inputs, forwardprop_aids)
2. 外側のGradientTapeで転置を計算:
   output_tangents = transpose_tape.gradient(grads, forwardprop_aids, nontrivial_tangents)
```

## データベース操作仕様

### 操作別データベース影響一覧

| 操作 | 対象テーブル | 操作種別 | 概要 |
|-----|-------------|---------|------|
| - | - | - | データベース操作なし |

### テーブル別操作詳細

データベース操作は発生しない。

## エラー処理

### エラーケース一覧

| エラーコード | エラー種別 | 発生条件 | 対処方法 |
|------------|----------|---------|---------|
| ValueError | primal重複 | 同一テンソルをprimalsに複数回指定 | 一意なテンソルを指定 |
| ValueError | tangent形状不一致 | tangentsがprimalsと異なる形状 | 同一形状を指定 |

### リトライ仕様

リトライは不要。

## トランザクション仕様

特になし。_TRACE_COUNT_CONSISTENCY_LOCKでトレースカウントの書き込み競合を防止。

## パフォーマンス要件

- 特殊ケース（Identity, ReadVariableOp）は即座に完了
- 一般ケースはGradientTape 2つを使うため、逆方向微分より遅い可能性
- `_jvp_function_cache` によるトレース結果のキャッシュでリトレースを削減
- `reduce_retracing=True` で形状緩和トレーシングを有効化（32回超過時）

## セキュリティ考慮事項

特になし。

## 備考

- 前方自動微分は中間活性を保持しないため、backward-over-forwardよりメモリ効率が良い
- ヘシアン-ベクトル積の計算はforward-over-backwardパターンが推奨

---

## コードリーディングガイド

本機能を理解するために参照すべきファイルと、推奨する読み解き順序を以下に示す。

### 推奨読解順序

#### Step 1: 特殊ケースJVPを理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 1-1 | forwardprop.py | `tensorflow/python/eager/forwardprop.py` | _SPECIAL_CASESとIdentity/ReadVariableOp |

**主要処理フロー**:
1. **39行目**: `_SPECIAL_CASES = {}` 辞書で特殊ケースを管理
2. **42-47行目**: `_identity_jvp` - Identity Opの直接JVP
3. **50行目**: `_SPECIAL_CASES["Identity"] = _identity_jvp`
4. **53-57行目**: `_read_variable_jvp` - ReadVariableOpの直接JVP
5. **60行目**: `_SPECIAL_CASES["ReadVariableOp"] = _read_variable_jvp`

#### Step 2: 一般的なJVP計算を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 2-1 | forwardprop.py | `tensorflow/python/eager/forwardprop.py` | _jvp_helper関数 |

**主要処理フロー**:
- **69-138行目**: `_jvp_helper` - メイン計算ロジック
- **87-90行目**: _TRACE_COUNTの更新
- **92-94行目**: 特殊ケースのチェック
- **101-113行目**: 内側のbackfunc_tapeで逆方向勾配を記録
- **126-131行目**: 外側のtranspose_tapeでJVP（転置）を計算
- **132-133行目**: output_tangentsの組み立て

**読解のコツ**: `forwardprop_aids` は `ones_like` テンソルで、逆方向勾配の転置を計算するためのダミー出力勾配として使用される。実際の計算に使われず、プルーニングされる。

#### Step 3: ディスパッチとキャッシュ

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 3-1 | forwardprop.py | `tensorflow/python/eager/forwardprop.py` | _jvp_dispatch関数 |

**主要処理フロー**:
- **187-193行目**: `_jvp_function_cache` と `TracingOptions` の設定
- **204行目**: `_TRACE_COUNT_LIMIT = 32`
- **207-223行目**: `_jvp_dispatch` - カウントに基づくexact/relaxedの選択
- **226行目**: `pywrap_tfe.TFE_Py_RegisterJVPFunction(_jvp_dispatch)` でC++に登録

#### Step 4: ForwardAccumulatorクラス

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 4-1 | forwardprop.py | `tensorflow/python/eager/forwardprop.py` | ForwardAccumulatorクラス |

**主要処理フロー**:
- **229行目**: `@tf_export("autodiff.ForwardAccumulator")`
- **340-349行目**: `__init__` - primals, tangentsの保存

### プログラム呼び出し階層図

```
tf.autodiff.ForwardAccumulator (forwardprop.py)
    |
    +-- __enter__ / __exit__
    |       +-- pywrap_tfe.TFE_Py_ForwardAccumulatorNew / Watch / SetRemove
    |
    +-- jvp(primal)
    |       +-- pywrap_tfe.TFE_Py_ForwardAccumulatorJVP
    |
    +-- [Op実行時] _jvp_dispatch (自動呼び出し)
            |
            +-- _jvp_helper_wrapper (TracingOptionsでトレース管理)
                |
                +-- _jvp_helper
                    |
                    +-- _SPECIAL_CASES.get(op_name) [特殊ケース]
                    |       +-- _identity_jvp / _read_variable_jvp
                    |
                    +-- [一般ケース]
                        +-- GradientTape (backfunc_tape) -> 逆方向勾配記録
                        +-- GradientTape (transpose_tape) -> 転置計算
```

### データフロー図

```
[入力]               [処理]                    [出力]

primals -----+
             |
tangents ----+--> ForwardAccumulator
                     |
                [Op実行] --> _jvp_dispatch
                     |            |
                     |     _jvp_helper
                     |       /        \
                     |  特殊ケース   一般ケース
                     |  (identity)  (backfunc + transpose)
                     |       \        /
                     |     output tangents
                     |
                jvp(primal) --> tangent Tensor
```

### 関連ファイル一覧

| ファイル | パス | 種別 | 役割 |
|---------|------|------|------|
| forwardprop.py | `tensorflow/python/eager/forwardprop.py` | ソース | ForwardAccumulator主実装 |
| forwardprop_util.py | `tensorflow/python/eager/forwardprop_util.py` | ソース | push_forwardprop_state等 |
| backprop.py | `tensorflow/python/eager/backprop.py` | ソース | GradientTape（転置計算用） |
| backprop_util.py | `tensorflow/python/eager/backprop_util.py` | ソース | IsTrainable判定 |
| pywrap_tfe.py | `tensorflow/python/pywrap_tfe.py` | バインディング | C++アキュムレータ管理 |
| function_cache.py | `tensorflow/core/function/polymorphism/function_cache.py` | ソース | JVP関数キャッシュ |
