rewriter.replaceOp(batch_norm, op->getResults());

    // Depending on the case, we may fuse the add, the relu, or both.
    if (!add_op || add_op.getZ().hasOneUse()) {
      // We fuse the Relu only if the add has a single use, otherwise we only
      // fuse the add itself.
      op->setAttr("activation_mode", rewriter.getStringAttr("Relu"));
      rewriter.replaceOp(relu_op, op->getResult(0));
    }

      // TODO(b/267700110): Handle multiple input/output cases.
      rewriter.replaceOp(replicated_input_op, replicated_input_op.getInputs());
    } else if (auto replicated_output_op =
                   dyn_cast_or_null<TF::TPUReplicatedOutputOp>(op)) {
      // TODO(b/267700110): Handle multiple input/output cases.
      rewriter.replaceOp(replicated_output_op, replicated_output_op.getInput());
    } else {
      return failure();

      auto tf_argreduce_op = rewriter.create<ArgReduce>(
          reduce_op.getLoc(), reduce_op->getResult(1).getType(), operand,
          reduction_indices);

      rewriter.replaceOp(reduce_op, {tf_reduce_op, tf_argreduce_op});
    } else {
      auto tf_reduce_op = rewriter.create<Reduce>(
          reduce_op.getLoc(), reduce_op->getResult(0).getType(), operand,
          reduction_indices,

        pad_op.getPaddingValue(), new_padding_low, new_padding_high,
        new_padding_interrier);

    Value orig_pad = CreateTranspose(rewriter, new_pad, transpose_perm);
    rewriter.replaceOp(pad_op, orig_pad);
    return success();
  }
};

// Transform reduce_window(transpose(x)) to transpose(reduce_window(x))
struct TransposeCommuteWithReduceWindow

          TFL::ConstBytesAttr::get(rewriter.getContext(), stringattr));
    }
  } else {
    tfl_custom.setCustomOptionAttr(
        TFL::ConstBytesAttr::get(rewriter.getContext(), ""));
  }

  rewriter.replaceOp(mhlo_custom_call, tfl_custom);
  return success();
}

std::optional<bool> IsCustomCallLegal(mhlo::CustomCallOp op) {
  if (op.getCallTargetName().starts_with("custom_call.")) {
    auto bc = op.getBackendConfig();

        num_parallel_calls_op.getY(), op.getDropRemainder(),
        batchInputOp.getF(), op.getOutputTypes(), op.getOutputShapes(),
        batchInputOp.getPreserveCardinality(), op.getMetadata());
    rewriter.replaceOp(op, {fused_op.getHandle()});
    return failure();
  }
};

#include "tensorflow/compiler/mlir/tensorflow/transforms/generated_tf_data_optimization.inc"
}  // namespace

        conv_op.getFeatureGroupCount(), conv_op.getBatchGroupCount(),
        conv_op.getPrecisionConfigAttr());
    // For static case, replace the convolution op now.
    if (!is_dynamic_broadcast) {
      rewriter.replaceOp(mul_op, {new_conv});
    } else {
      // For dynamic case, create new shape_of op and replace uses.
      shape_of_op =
          dyn_cast_or_null<mhlo::DynamicBroadcastInDimOp>(bcast_or_const_op)

    // compilers, so avoid removing these Identity ops.
    if (user->hasTrait<OpTrait::IsTerminator>()) {
      return failure();
    }
  }

  rewriter.replaceOp(identity, identity.getInput());
  return success();
}

}  // namespace quant

          DenseFPElementsAttr::get(
              mlir::dyn_cast<ShapedType>(const_op.getValue().getType())
                  .clone(rewriter.getBF16Type()),
              bfloat16_values));
    }

    rewriter.replaceOp(op, rewriter.create(state)->getResults());

    return success();
  }
};

class BitcastConvertOpPattern
    : public OpConversionPattern<mlir::stablehlo::BitcastConvertOp> {
 public:

    createScfThenOrElse(else_region, scf_if_op.getElseRegion(),
                        tf_if_region_return_type, rewriter);

    // Replace the `tf.IfRegion` op with the results of the `scf.if` op.
    rewriter.replaceOp(op, scf_if_op.getResults());
    return success();
  }
};

/// Convert the `tf.WhileRegion` op to the `scf.while` op.
class ConvertWhileRegionOp : public OpRewritePattern<WhileRegionOp> {
 public: