// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: tfr-opt %s -tfr-decompose -tfr-raise-to-tf -canonicalize -verify-diagnostics -split-input-file | FileCheck %s

//=================> User models, from GraphDef <====================

// CHECK-LABEL: my_identity
func.func @my_identity(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
  %0 = "tf.MyIdentity"(%arg0) : (tensor<2x3xf32>) -> tensor<2x3xf32>
  func.return %0 : tensor<2x3xf32>

// CHECK-NEXT: return %arg0 : tensor<2x3xf32>
}

// CHECK-LABEL: my_rsqrt
func.func @my_rsqrt(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
  %0 = "tf.MyRsqrt"(%arg0) : (tensor<2x3xf32>) -> tensor<3x2x3xf32>
  func.return %0 : tensor<3x2x3xf32>

// CHECK-NEXT: %[[RE:.*]] = "tf.RiscReciprocal"(%arg0) : (tensor<2x3xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[SQRT:.*]] = "tf.RiscSqrt"(%[[RE]]) : (tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[SQRT]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}

// CHECK-LABEL: my_leaky_relu
func.func @my_leaky_relu(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
  %0 = "tf.MyLeakyRelu"(%arg0) {alpha=3.0 : f32} : (tensor<2x3xf32>) -> tensor<3x2x3xf32>
  func.return %0 : tensor<3x2x3xf32>

// CHECK-NEXT: %[[ALPHA:.*]] = "tf.Const"() <{value = dense<3.000000e+00> : tensor<f32>}> : () -> tensor<f32>
// CHECK-NEXT: %[[SHAPE:.*]] = "tf.RiscShape"(%arg0) {T = i32} : (tensor<2x3xf32>) -> tensor<*xi32>
// CHECK-NEXT: %[[ALPHA1:.*]] = "tf.RiscBroadcast"(%[[ALPHA]], %[[SHAPE]]) : (tensor<f32>, tensor<*xi32>) -> tensor<*xf32>
// CHECK-NEXT: %[[MAX:.*]] = "tf.RiscMaximum"(%arg0, %[[ALPHA1]]) : (tensor<2x3xf32>, tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[MAX]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}

// CHECK-LABEL: my_leaky_relu_with_default
func.func @my_leaky_relu_with_default(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
  %0 = "tf.MyLeakyRelu"(%arg0) : (tensor<2x3xf32>) -> tensor<3x2x3xf32>
  func.return %0 : tensor<3x2x3xf32>

// CHECK-NEXT: %[[ALPHA:.*]] = "tf.Const"() <{value = dense<2.000000e-01> : tensor<f32>}> : () -> tensor<f32>
// CHECK-NEXT: %[[SHAPE:.*]] = "tf.RiscShape"(%arg0) {T = i32} : (tensor<2x3xf32>) -> tensor<*xi32>
// CHECK-NEXT: %[[ALPHA1:.*]] = "tf.RiscBroadcast"(%[[ALPHA]], %[[SHAPE]]) : (tensor<f32>, tensor<*xi32>) -> tensor<*xf32>
// CHECK-NEXT: %[[MAX:.*]] = "tf.RiscMaximum"(%arg0, %[[ALPHA1]]) : (tensor<2x3xf32>, tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[MAX]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}

// CHECK-LABEL: my_cast
func.func @my_cast(%arg0: tensor<2x3xf32>) -> tensor<2x3xi32> {
  %0 = "tf.MyCast"(%arg0) {Tout=i32} : (tensor<2x3xf32>) -> tensor<2x3xi32>
  func.return %0 : tensor<2x3xi32>

// CHECK-NEXT: %[[CAST:.*]] = "tf.RiscCast"(%arg0) {Tout = i32} : (tensor<2x3xf32>) -> tensor<*xi32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[CAST]]) <{shape = #tf_type.shape<2x3>}> : (tensor<*xi32>) -> tensor<2x3xi32>
// CHECK-NEXT: return %[[ES]] : tensor<2x3xi32>
}

// CHECK-LABEL: my_pack_single_input
func.func @my_pack_single_input(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
  %0 = "tf.MyPack"(%arg0) {N=1:i32, axis=0:i32} : (tensor<2x3xf32>) -> tensor<3x2x3xf32>
  func.return %0 : tensor<3x2x3xf32>

// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT: %[[ED:.*]] = "tf.ExpandDims"(%arg0, %[[AXIS]]) : (tensor<2x3xf32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[ED]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}

// CHECK-LABEL: my_pack_multiple_inputs
func.func @my_pack_multiple_inputs(%arg0: tensor<2x3xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
  %0 = "tf.MyPack"(%arg0, %arg1, %arg2) {N=3:i32, axis=0:i32} : (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<3x2x3xf32>
  func.return %0 : tensor<3x2x3xf32>

// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT: %[[ED0:.*]] = "tf.ExpandDims"(%arg0, %[[AXIS]]) : (tensor<2x3xf32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ED1:.*]] = "tf.ExpandDims"(%arg1, %[[AXIS]]) : (tensor<2x3xf32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: %[[CC0:.*]] = "tf.RiscConcat"(%[[ED0]], %[[ED1]]) {axis = 0 : i32} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ED2:.*]] = "tf.ExpandDims"(%arg2, %[[AXIS]]) : (tensor<2x3xf32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: %[[CC1:.*]] = "tf.RiscConcat"(%[[CC0]], %[[ED2]]) {axis = 0 : i32} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[CC1]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}

// CHECK-LABEL: my_add_n_single_input
func.func @my_add_n_single_input(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
  %0 = "tf.MyAddN"(%arg0) {N=1:i32} : (tensor<2x3xf32>) -> tensor<2x3xf32>
  func.return %0 : tensor<2x3xf32>

// CHECK-NEXT: return %arg0 : tensor<2x3xf32>
}

// CHECK-LABEL: my_add_n_multiple_inputs
func.func @my_add_n_multiple_inputs(%arg0: tensor<2x3xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<2x3xf32>) -> tensor<2x3xf32> {
  %0 = "tf.MyAddN"(%arg0, %arg1, %arg2) {N=3:i32} : (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
  func.return %0 : tensor<2x3xf32>

// CHECK-NEXT: %[[ADD0:.*]] = "tf.RiscAdd"(%arg0, %arg1) : (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ADD1:.*]] = "tf.RiscAdd"(%[[ADD0]], %arg2) : (tensor<*xf32>, tensor<2x3xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[ADD1]]) <{shape = #tf_type.shape<2x3>}> : (tensor<*xf32>) -> tensor<2x3xf32>
// CHECK-NEXT: return %[[ES]] : tensor<2x3xf32>
}

// CHECK-LABEL: my_map_and_batch_dataset
func.func @my_map_and_batch_dataset(%input: tensor<*x!tf_type.variant>,
                               %other1: tensor<*xf32>,
                               %other2: tensor<*xi32>) -> tensor<*x!tf_type.variant> {
  %0 = "tf.MyMapAndBatchDataset"(%input, %other1, %other2)
    {batch_size=1000 : i64, num_parallel_calls = 8 : i64, drop_remainder = 0 : i1,
     func = @"__some_func", output_types = [f32], output_shapes = [#tf_type.shape<>], preserve_cardinality = true}
    : (tensor<*x!tf_type.variant>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.variant>
  func.return %0 : tensor<*x!tf_type.variant>

// CHECK-DAG: %[[BATCH:.*]] = "tf.Const"() <{value = dense<1000> : tensor<i64>}> : () -> tensor<i64>
// CHECK-DAG: %[[PARAL:.*]] = "tf.Const"() <{value = dense<8> : tensor<i64>}> : () -> tensor<i64>
// CHECK-DAG: %[[KEEP:.*]] = "tf.Const"() <{value = dense<false> : tensor<i1>}> : () -> tensor<i1>
// CHECK: %[[CAST:.*]] = "tf.Cast"(%arg2) <{Truncate = false}> : (tensor<*xi32>) -> tensor<*xf32>
// CHECK: %[[RET:.*]] = "tf.MapAndBatchDatasetV0"(%arg0, %[[BATCH]], %[[PARAL]], %[[KEEP]], %arg1, %[[CAST]])
// CHECK-SAME: {f = @__some_func, output_shapes = [#tf_type.shape<>], output_types = [f32], preserve_cardinality = true} : (tensor<*x!tf_type.variant>, tensor<i64>, tensor<i64>, tensor<i1>, tensor<*xf32>, tensor<*xf32>) -> tensor<*x!tf_type.variant>
// CHECK: return %[[RET]] : tensor<*x!tf_type.variant>
}

//=================> decomposition functions, translated from tf.compose api <====================
tfr.func @tf__my_identity(%value: !tfr.tensor) -> !tfr.tensor {
  tfr.return %value : !tfr.tensor
}

tfr.func @tf__my_cast(%value: !tfr.tensor, %tout: !tfr.attr{tfr.name="Tout"}) -> !tfr.tensor {
  %0 = tfr.call @tf__risc_cast(%value, %tout) : (!tfr.tensor, !tfr.attr) -> !tfr.tensor
  tfr.return %0 : !tfr.tensor
}

tfr.func @tf__my_rsqrt(%value: !tfr.tensor) -> !tfr.tensor {
  %1 = tfr.call @tf__risc_reciprocal(%value) : (!tfr.tensor) -> !tfr.tensor
  %2 = tfr.call @tf__risc_sqrt(%1) : (!tfr.tensor) -> !tfr.tensor
  tfr.return %2 : !tfr.tensor
}

tfr.func @tf__my_leaky_relu(%value: !tfr.tensor, %alpha: f32 {tfr.name="alpha", tfr.default=0.2:f32}) -> !tfr.tensor {
  %1 = tfr.call @tf__risc_shape(%value) : (!tfr.tensor) -> !tfr.tensor
  %2 = "tfr.constant_tensor"(%alpha) : (f32) -> tensor<f32>
  %t = "tfr.cast"(%2) : (tensor<f32>) -> !tfr.tensor
  %3 = tfr.call @tf__risc_broadcast(%t, %1) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
  %4 = tfr.call @tf__risc_maximum(%value, %3) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
  tfr.return %4  : !tfr.tensor
}

// TODO(fengliuai): use shape dialect to manipulate the shape then this can be decomposed further.
tfr.func @tf__my_expand_dims(%value: !tfr.tensor, %axis: i32 {tfr.name="axis"}) -> !tfr.tensor {
  %axis_cst = "tfr.constant_tensor"(%axis) : (i32) -> tensor<i32>
  %dim = "tfr.cast"(%axis_cst) : (tensor<i32>) -> !tfr.tensor
  %0 = tfr.call @tf__expand_dims(%value, %dim) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
  tfr.return %0 : !tfr.tensor
}

tfr.func @tf__my_pack(%values: !tfr.tensor_list,
                      %n: i32 {tfr.name="N"},
                      %axis: i32 {tfr.name="axis"}) -> !tfr.tensor {
  %index = arith.constant 0 : index
  %cst = arith.constant 1 : i32
  %eq = arith.cmpi eq, %n, %cst : i32
  %v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
  %temp = tfr.call @tf__my_expand_dims(%v1, %axis) : (!tfr.tensor, i32) -> !tfr.tensor
  %res = scf.if %eq -> !tfr.tensor {
    scf.yield %temp : !tfr.tensor
  } else {
    %step = arith.index_cast %cst : i32 to index
    %end = arith.index_cast %n : i32 to index
    %reduce = scf.for %i = %step to %end step %step iter_args(%reduce_iter=%temp) -> !tfr.tensor {
      %v = tfr.get_element %values[%i] : (!tfr.tensor_list, index) -> !tfr.tensor
      %temp1 =  tfr.call @tf__my_expand_dims(%v, %axis) : (!tfr.tensor, i32) -> !tfr.tensor
      %reduce_next =  tfr.call @tf__risc_concat(%reduce_iter, %temp1, %axis) : (!tfr.tensor, !tfr.tensor, i32) -> !tfr.tensor
      scf.yield %reduce_next : !tfr.tensor
    }
    scf.yield %reduce : !tfr.tensor
  }
  tfr.return %res : !tfr.tensor
}

tfr.func @tf__my_add_n(%values: !tfr.tensor_list,
                       %n: i32 {tfr.name="N"}) -> !tfr.tensor {
  %index = arith.constant 0 : index
  %cst = arith.constant 1 : i32
  %eq = arith.cmpi eq, %n, %cst : i32
  %v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
  %res = scf.if %eq -> !tfr.tensor {
    scf.yield %v1 : !tfr.tensor
  } else {
    %step = arith.index_cast %cst : i32 to index
    %end = arith.index_cast %n : i32 to index
    %reduce = scf.for %i = %step to %end step %step iter_args(%reduce_iter=%v1) -> !tfr.tensor {
      %v = tfr.get_element %values[%i] : (!tfr.tensor_list, index) -> !tfr.tensor
      %reduce_next =  tfr.call @tf__risc_add(%reduce_iter, %v) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor
      scf.yield %reduce_next : !tfr.tensor
    }
    scf.yield %reduce : !tfr.tensor
  }
  tfr.return %res : !tfr.tensor
}

tfr.func @tf__my_map_and_batch_dataset(
    %input_dataset: !tfr.tensor,
    %other_arguments: !tfr.tensor_list,
    %batch_size: i64 {tfr.name="batch_size"},
    %num_parallel_calls: i64 {tfr.name="num_parallel_calls"},
    %drop_remainder: i1 {tfr.name="drop_remainder"},
    %f: !tfr.attr {tfr.name="func"},
    %output_types: !tfr.attr {tfr.name="output_types"},
    %output_shapes: !tfr.attr {tfr.name="output_shapes"},
    %preserve_cardinality: i1 {tfr.name="preserve_cardinality", tfr.default=false}) -> !tfr.tensor {
  %batch = "tfr.constant_tensor"(%batch_size) : (i64) -> tensor<i64>
  %batch1 = "tfr.cast"(%batch) : (tensor<i64>) -> !tfr.tensor
  %calls = "tfr.constant_tensor"(%num_parallel_calls) : (i64) -> tensor<i64>
  %calls1 = "tfr.cast"(%calls) : (tensor<i64>) -> !tfr.tensor
  %drop = "tfr.constant_tensor"(%drop_remainder) : (i1) -> tensor<i1>
  %drop1 = "tfr.cast"(%drop) : (tensor<i1>) -> !tfr.tensor
  %ret = tfr.call @tf__map_and_batch_dataset_v0(%input_dataset, %batch1, %calls1, %drop1, %other_arguments, %f, %output_types, %output_shapes, %preserve_cardinality)
    : (!tfr.tensor, !tfr.tensor, !tfr.tensor, !tfr.tensor, !tfr.tensor_list, !tfr.attr, !tfr.attr, !tfr.attr, i1) -> !tfr.tensor
  tfr.return %ret : !tfr.tensor
}

//=================>  signatures of the primitive ops with kernels, modeled as external TFR function <==
tfr.func @tf__risc_cast_(!tfr.tensor, !tfr.attr{tfr.name="Tout"}) -> !tfr.tensor<Tout> attributes{Tout}
tfr.func @tf__risc_add_(!tfr.tensor<T>, !tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_concat_(!tfr.tensor<T>, !tfr.tensor<T>, i32{tfr.name="axis"}) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_broadcast_(!tfr.tensor<T>, !tfr.tensor<Tidx>) -> !tfr.tensor<T> attributes{T, Tidx}
tfr.func @tf__risc_reciprocal_(!tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_sqrt_(!tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_shape_(!tfr.tensor, !tfr.attr{tfr.name="T", tfr.default=i32}) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__risc_maximum_(!tfr.tensor<T>, !tfr.tensor<T>) -> !tfr.tensor<T> attributes{T}
tfr.func @tf__expand_dims_(!tfr.tensor<T>, !tfr.tensor<Tdim>) -> !tfr.tensor<T> attributes{T, Tdim}
tfr.func @tf__map_and_batch_dataset_v0_(!tfr.tensor<T>, !tfr.tensor, !tfr.tensor, !tfr.tensor, !tfr.tensor_list<Targuments>,
  !tfr.attr{tfr.name="f"}, !tfr.attr{tfr.name="output_types"}, !tfr.attr{tfr.name="output_shapes"}, i1{tfr.name="preserve_cardinality"})
  -> !tfr.tensor<T> attributes{T, Targuments}