# 機能設計書 95-確率的キャスト

## 概要

本ドキュメントは、TensorFlowにおける確率的キャスト（Stochastic Cast）機能の設計を記述する。確率的な型変換（低精度化）を提供する。

### 本機能の処理概要

**業務上の目的・背景**：ニューラルネットワークの量子化において、通常の丸め（rounding）ではバイアスが生じる場合がある。確率的キャスト（Stochastic Rounding）は、浮動小数点値を確率的に丸めることで、期待値として正確な値を保持する手法である。これにより、低精度（INT8/INT16/INT32）への型変換時の精度損失を最小限に抑える。

**機能の利用シーン**：混合精度学習における低精度への変換、量子化対応学習（QAT）でのストカスティックラウンディング、勾配の低精度化による通信効率向上（分散学習）などに使用される。

**主要な処理内容**：
1. `StochasticCastToInt` - 浮動小数点テンソルを確率的丸めで整数型に変換
2. Philox乱数生成器を使用した確率的丸め値の決定
3. 入力テンソルと同じ形状の出力を生成

**関連システム・外部連携**：乱数生成フレームワーク（RNG）と連携。Philoxアルゴリズムによるカウンターベース乱数を使用。

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

## 関連画面

| 画面No | 画面名 | 関連種別 | 関連する操作・処理 |
|--------|--------|----------|------------------|
| - | - | - | 本機能はAPI機能であり、直接関連する画面はない |

## 機能種別

計算処理（確率的型変換）

## 入力仕様

### 入力パラメータ

| パラメータ名 | 型 | 必須 | 説明 | バリデーション |
|-------------|-----|-----|------|---------------|
| input | Tin (half/bfloat16/float32/float64) | Yes | 変換対象の浮動小数点テンソル | - |
| key | uint64 | Yes | RNGキー（Philox） | rank 1, サイズ = RNG_KEY_SIZE |
| counter | uint64 | Yes | RNGカウンター | rank 1 |
| alg | int32 | Yes | RNGアルゴリズム指定 | スカラー |
| Tout | type | Yes（属性） | 出力型: int8, int16, int32 | - |

### 入力データソース

計算グラフ内のテンソルおよびRNG状態。

## 出力仕様

### 出力データ

| 項目名 | 型 | 説明 |
|--------|-----|------|
| output | Tout (int8/int16/int32) | 確率的に丸められた整数テンソル（入力と同じ形状） |

### 出力先

計算グラフ内の後続操作への入力。

## 処理フロー

### 処理シーケンス

```
1. 入力形状の検証
   └─ key: rank 1, サイズ = RNG_KEY_SIZE
   └─ counter: rank 1
   └─ alg: スカラー
2. 乱数生成
   └─ Philox RNGで確率的丸め用のランダム値を生成
3. 確率的丸めの実行
   └─ 各要素について: floor(x) + bernoulli(x - floor(x))
4. 出力型への変換
   └─ 結果をTout型（int8/int16/int32）にキャスト
```

### フローチャート

```mermaid
flowchart TD
    A[入力: float Tensor] --> B[RNG状態検証]
    B --> C[Philox乱数生成]
    C --> D[確率的丸め計算]
    D --> E[整数型キャスト]
    E --> F[出力: int Tensor]
```

## ビジネスルール

### 業務ルール

| ルールNo | ルール名 | 内容 | 適用条件 |
|---------|---------|------|---------|
| BR-95-01 | 期待値保存 | 確率的丸めの期待値は元の浮動小数点値と一致する | 常時 |
| BR-95-02 | key形状制約 | keyのサイズはRNG_KEY_SIZEと一致する必要がある | 常時 |
| BR-95-03 | 出力形状保存 | 出力テンソルの形状は入力テンソルと同一 | 常時 |

### 計算ロジック

確率的丸め（Stochastic Rounding）：
- `x`を確率的に`floor(x)`または`ceil(x)`に丸める
- `P(output = ceil(x)) = x - floor(x)`
- `P(output = floor(x)) = ceil(x) - x`
- 期待値: `E[output] = x`

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

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

本機能はデータベース操作を行わない。

## エラー処理

### エラーケース一覧

| エラーコード | エラー種別 | 発生条件 | 対処方法 |
|------------|----------|---------|---------|
| InvalidArgument | 形状不正 | keyのサイズがRNG_KEY_SIZEでない | 正しいサイズのkeyを使用 |
| InvalidArgument | rank不正 | key, counterがrank 1でない、algがスカラーでない | 正しいrankを使用 |

### リトライ仕様

リトライ不要。

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

該当なし。入力のRNG状態を除きステートレス。

## パフォーマンス要件

- 通常のキャスト操作と比較して、乱数生成分の追加コストがある
- Philox RNGは高速なカウンターベース生成器であるため、オーバーヘッドは限定的

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

- 確率的丸めにPhilox RNGを使用するが、暗号学的安全性は保証しない

## 備考

- TODO(b/232442915): 浮動小数点から低精度浮動小数点への丸めサポートが将来的に追加予定
- 入力型: half, bfloat16, float32, float64
- 出力型: int8, int16, int32

---

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

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

### 推奨読解順序

#### Step 1: Op定義を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 1-1 | stochastic_cast_op.cc | `tensorflow/core/ops/stochastic_cast_op.cc` | StochasticCastToInt Opの定義 |

**主要処理フロー**:
1. **27-46行目**: `StochasticCastToInt` Op登録
2. **33行目**: Tin属性 - half, bfloat16, float32, float64をサポート
3. **34行目**: Tout属性 - int8, int16, int32をサポート
4. **35-46行目**: 形状推論ラムダ
5. **38行目**: key形状検証 - rank 1
6. **39行目**: counter形状検証 - rank 1
7. **40行目**: alg形状検証 - スカラー
8. **42-43行目**: keyサイズ検証 - RNG_KEY_SIZEと一致
9. **44行目**: 出力形状 = 入力形状

**読解のコツ**: `rng_alg.h`の`RNG_KEY_SIZE`定数を参照。Philox RNGのkey/counter/algの3つの状態要素の意味を理解する。

#### Step 2: RNGフレームワークを理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 2-1 | rng_alg.h | `tensorflow/core/framework/rng_alg.h` | RNG_KEY_SIZE定数と乱数アルゴリズム定義 |

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

```
StochasticCastToInt Op
    │
    ├─ 形状推論
    │      ├─ WithRank(key, 1)
    │      ├─ WithRank(counter, 1)
    │      ├─ WithRank(alg, 0)
    │      ├─ WithValue(key[0], RNG_KEY_SIZE)
    │      └─ output_shape = input_shape
    │
    └─ カーネル実行
           ├─ Philox RNG初期化 (key, counter, alg)
           ├─ 確率的丸め値生成
           └─ 整数型キャスト
```

### データフロー図

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

float Tensor        ───▶  StochasticCastToInt       ───▶  int Tensor
key [RNG_KEY_SIZE]  ───▶   ├─ Philox RNG生成               (同形状)
counter [N]         ───▶   ├─ 確率的丸め
alg (scalar)        ───▶   └─ 型変換
```

### 関連ファイル一覧

| ファイル | パス | 種別 | 役割 |
|---------|------|------|------|
| stochastic_cast_op.cc | `tensorflow/core/ops/stochastic_cast_op.cc` | ソース | Op定義と形状推論 |
| rng_alg.h | `tensorflow/core/framework/rng_alg.h` | ヘッダ | RNG_KEY_SIZE定数定義 |
