tree 4ad847aac301d2e6a60dcb794b1a1a588e34373d
parent cc0626cc42b2cb1542c97e3aa450b8aa3ce303c7
author Yuanqiang Liu <liuyuanqiang.yqliu@bytedance.com> 1657213065 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> 1657213302 -0700

PR #56483: [MHLO] unfuse batch norm ops with shape dialect, and split inference …

Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/56483

…and training pattern

1. lowering with `shape` dialect. The original lowering pattern is not a Progressive Lowering. As the `shape` dialect is more mature in MLIR, I think lowering dynamic shape pattern with `shape` dialect is more suitable.
2. split inference and training pattern for flexible integration. The original patterns have both `mhlo.batch_norm_inference` and `mhlo.batch_norm_training` lowering pattern. Sometimes we want to only lower one of them, so I split these pattern.
Copybara import of the project:

--
f2bbaa34d021cf7c66d18d10d0f0a2535bceb54e by liuyuanqiang <liuyuanqiang.yqliu@bytedance.com>:

[MHLO] unfuse batch norm ops with shape dialect, and split inference and training pattern

Merging this change closes #56483

COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/tensorflow/pull/56483 from qingyunqu:fix-unfuse-batch-norm f2bbaa34d021cf7c66d18d10d0f0a2535bceb54e
PiperOrigin-RevId: 459537687