text: # Licensed to the Apache Software Foundation (ASF) under one
  # or more contributor license agreements.  See the NOTICE file
  # distributed with this work for additional information
  # regarding copyright ownership.  The ASF licenses this file
  # to you 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.
  """ Support level3 operator test cases.
  """
  import numpy as np
  import pytest
  import tvm
  from tvm import te
  from tvm import relay
  from tvm.relay import create_executor, transform
  from tvm.relay.testing import check_grad, run_infer_type
  import tvm.testing
  
  
  def verify_func(func, data, ref_res, target_device=tvm.testing.enabled_targets()):
      assert isinstance(data, list)
      for target, dev in target_device:
          for kind in ["vm", "debug"]:
              mod = tvm.ir.IRModule.from_expr(func)
              op_res = relay.create_executor(kind, mod=mod, device=dev, target=target).evaluate()(
                  *data
              )
              if isinstance(op_res, tvm.runtime.container.ADT):
                  assert len(op_res) == len(
                      ref_res
                  ), "Outputs from TVM and Python implementation must be equal "
                  for op_result, ref_result in zip(op_res, ref_res):
                      tvm.testing.assert_allclose(op_result.numpy(), ref_result, rtol=1e-5)
              else:
                  tvm.testing.assert_allclose(op_res.numpy(), ref_res, rtol=1e-5)
              relay.backend.compile_engine.get().clear()
  
  
  @tvm.testing.uses_gpu
  def test_dyn_reshape():
      def verify_reshape(shape, newshape, oshape):
          x = relay.var("x", relay.TensorType(shape, "float32"))
          y = relay.var("y", relay.TensorType((len(newshape),), "int64"))
          z = relay.reshape(x, y)
  
          func = relay.Function([x, y], z)
          x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32")
          x_data = np.ones(shape).astype("float32")
          ref_res = np.reshape(x_data, oshape)
          check_grad(
              run_infer_type(func),
              inputs=[x_data, np.array(newshape).astype("int64")],
              test_inputs=[x_data],
              eps=1e-3,
          )
          verify_func(func, [x_data, np.array(newshape).astype("int64")], ref_res)
  
      verify_reshape((2, 3, 4), (8, 3), (8, 3))
      verify_reshape((4, 7), (2, 7, 2), (2, 7, 2))
      verify_reshape((2, 3, 4), (4, 0, 2), (4, 3, 2))
      verify_reshape((2, 3, 4), (2, 0, 0), (2, 3, 4))
      verify_reshape((2, 3, 4), (0, -1), (2, 12))
      verify_reshape((2, 3, 4), (-1, 0), (8, 3))
      verify_reshape((2, 3, 4), (-3, 4), (6, 4))
      verify_reshape((2, 3, 4, 5), (-3, -3), (6, 20))
      verify_reshape((2, 3, 4), (0, -3), (2, 12))
  
  
  @tvm.testing.uses_gpu
  def test_dyn_shape_reshape():
      def verify_reshape(shape, newshape, oshape):
          x = relay.var("x", relay.TensorType(shape, "float32"))
          y = relay.var("y", relay.TensorType(newshape, "float32"))
          z = relay.reshape(x, relay.shape_of(y))
  
          func = relay.Function([x, y], z)
          x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32")
          y_data = np.random.uniform(low=-1, high=1, size=newshape).astype("float32")
          ref_res = np.reshape(x_data, oshape)
          check_grad(run_infer_type(func), inputs=[x_data, y_data], eps=1e-3)
          verify_func(func, [x_data, y_data], ref_res)
  
      verify_reshape((2, 3, 4), (8, 3), (8, 3))
      verify_reshape((4, 7), (2, 7, 2), (2, 7, 2))
  
  
  @tvm.testing.uses_gpu
  def test_dyn_tile():
      def verify_tile(dshape, reps):
          x = relay.var("x", relay.TensorType(dshape, "float32"))
          r = relay.var("reps", relay.TensorType((len(reps),), "float32"))
          z = relay.tile(x, r)
  
          func = relay.Function([x, r], z)
          x_data = np.random.uniform(low=-1, high=1, size=dshape).astype("float32")
          ref_res = np.tile(x_data, reps=reps)
          reps_data = np.array(reps).astype("float32")
          verify_func(func, [x_data, np.array(reps).astype("float32")], ref_res)
  
      verify_tile((2, 3, 4), (3, 2, 1))
      verify_tile((2, 3, 4), (1, 2))
      verify_tile((2, 3), (3, 2, 1))
  
  
  @tvm.testing.uses_gpu
  def test_dyn_zeros_ones():
      def verify_zeros_ones(shape, dtype):
          for op, ref in [(relay.zeros, np.zeros), (relay.ones, np.ones)]:
              rank = len(shape)
              dyn_shape = relay.Var("shape", relay.ty.TensorType((rank,), "int64"))
              y = op(dyn_shape, dtype)
              yy = run_infer_type(y)
              assert yy.checked_type == relay.ty.TensorType((relay.Any(),) * rank, dtype)
  
              func = relay.Function([dyn_shape], y)
              ref_res = ref(shape, dtype)
              verify_func(func, [np.array(shape).astype("int64")], ref_res.astype("int64"))
  
      verify_zeros_ones((1, 3), "int64")
      verify_zeros_ones((8, 9, 1, 2), "float32")
  
  
  @tvm.testing.uses_gpu
  def test_dyn_full():
      def verify_full(fill_value, src_shape, dtype):
          x = relay.var("x", relay.scalar_type(dtype))
          rank = len(src_shape)
          dyn_src_shape = relay.var("dyn_scr_shape", relay.ty.TensorType((rank,), "int64"))
          z = relay.full(x, dyn_src_shape, dtype)
          func = relay.Function([x, dyn_src_shape], z)
          ref_res = np.full(src_shape, fill_value).astype(dtype)
  
          verify_func(
              func, [np.array(fill_value).astype(dtype), np.array(src_shape).astype("int64")], ref_res
          )
  
      verify_full(4, (1, 3, 4, 4), "int32")
      verify_full(4, (1, 3, 4, 4), "int64")
      verify_full(4.0, (2, 50), "float32")
  
  
  @tvm.testing.uses_gpu
  def test_dyn_sparse_to_dense():
      def verify_sparse_to_dense(sparse_indices, sparse_values, default_value, output_shape, xpected):
          sparse_indices_data = np.array(sparse_indices)
          sparse_values_data = np.array(sparse_values)
          default_value_data = np.array(default_value)
          output_shape_data = np.array(output_shape)
  
          a = relay.var(
              "a", relay.TensorType(sparse_indices_data.shape, str(sparse_indices_data.dtype))
          )
          b = relay.var(
              "b", relay.TensorType(sparse_values_data.shape, str(sparse_values_data.dtype))
          )
          output_shape_var = relay.var(
              "output_shape", relay.TensorType(output_shape_data.shape, str(output_shape_data.dtype))
          )
          if default_value is None:
              args = [a, b, output_shape_var]
              d = relay.sparse_to_dense(a, output_shape_var, b)
          else:
              c = relay.var(
                  "c", relay.TensorType(default_value_data.shape, str(default_value_data.dtype))
              )
              args = [a, b, c, output_shape_var]
              d = relay.sparse_to_dense(a, output_shape_var, b, c)
  
          zz = run_infer_type(d)
          assert len(zz.checked_type.shape) == len(output_shape)
  
          func = relay.Function(args, d)
  
          if default_value is None:
              arguments = [sparse_indices_data, sparse_values_data, output_shape_data]
          else:
              arguments = [
                  sparse_indices_data,
                  sparse_values_data,
                  default_value_data,
                  output_shape_data,
              ]
  
          verify_func(func, arguments, xpected)
  
      verify_sparse_to_dense(1, 3, 0, [5], [0, 3, 0, 0, 0])  # scalar
      verify_sparse_to_dense([0, 1, 4], [3, 3, 3], 0, [5], [3, 3, 0, 0, 3])  # vector
      verify_sparse_to_dense(
          [[0, 0], [1, 2]], [1, 2], 0, [3, 4], [[1, 0, 0, 0], [0, 0, 2, 0], [0, 0, 0, 0]]
      )  # nXd
      verify_sparse_to_dense(
          [[0, 0, 0], [1, 2, 3]],
          [1, 2],
          4,
          [2, 3, 4],
          [[[1, 4, 4, 4], [4, 4, 4, 4], [4, 4, 4, 4]], [[4, 4, 4, 4], [4, 4, 4, 4], [4, 4, 4, 2]]],
      )  # nXd
      verify_sparse_to_dense(
          [0, 1, 4], [3.1, 3.1, 3.1], 3.5, [5], [3.1, 3.1, 3.5, 3.5, 3.1]
      )  # floats
      verify_sparse_to_dense(1, 3, None, [5], [0, 3, 0, 0, 0])  # default value not specified
  
  
  @pytest.mark.parametrize(
      "sparse_indices, sparse_values, dense_shape, default_value",
      [
          (
              np.array([[0, 1], [0, 3], [2, 0], [3, 1]], dtype=np.int64),
              np.array([1, 2, 3, 4], dtype=np.int64),
              np.array([5, 6], dtype=np.int64),
              np.array([10], dtype=np.int64),
          ),
          (
              np.array([[1, 1, 1], [1, 3, 1], [2, 0, 5], [3, 1, 6]], dtype=np.int64),
              np.array([1, 2, 3, 4], dtype=np.int64),
              np.array([7, 7, 7], dtype=np.int64),
              np.array([5], dtype=np.int64),
          ),
          (
              np.array([[1], [2]], dtype=np.int64),
              np.array([7, 8], dtype=np.int64),
              np.array([5], dtype=np.int64),
              np.array([4], dtype=np.int64),
          ),
          (
              np.ones((0, 1), dtype=np.int64),
              np.array([], dtype=np.int64),
              np.array([5], dtype=np.int64),
              np.array([4], dtype=np.int64),
          ),
          (
              np.ones((0, 3), dtype=np.int64),
              np.array([], dtype=np.int64),
              np.array([9, 3, 7], dtype=np.int64),
              np.array([100], dtype=np.int64),
          ),
      ],
  )
  @pytest.mark.parametrize("dtype", [np.int64, np.int32])
  @pytest.mark.parametrize("use_dyn", [True, False])
  def test_sparse_fill_empty_rows(
      sparse_indices, sparse_values, dense_shape, default_value, dtype, use_dyn
  ):
      def ref_sparse_fill_empty_rows(
          sparse_indices: np.ndarray,
          sparse_values: np.ndarray,
          dense_shape: np.ndarray,
          default_value: np.ndarray,
      ) -> None:
          """
          This function calculates the expected output of sparse_fill_empty_rows operator given the
          inputs.
          """
  
          def check_add_rows(current_idx, limit_idx):
              while current_idx < limit_idx:
                  new_sparse_indices.append([current_idx] + [0] * (num_cols - 1))
                  new_sparse_values.append(default_value[0])
                  empty_row_indicator[current_idx] = True
                  current_idx += 1
  
              return current_idx
  
          current_idx = 0
          new_sparse_indices = []
          new_sparse_values = []
          empty_row_indicator = [False for _ in range(dense_shape[0])]
          num_cols = sparse_indices.shape[1]
          for sparse_row, sparse_value in zip(sparse_indices, sparse_values):
              limit_idx = sparse_row[0]
              current_idx = check_add_rows(current_idx, limit_idx)
              new_sparse_indices.append(list(sparse_row))
              new_sparse_values.append(sparse_value)
              current_idx = limit_idx + 1
  
          check_add_rows(current_idx, dense_shape[0])
          return new_sparse_indices, new_sparse_values, empty_row_indicator
  
      def verify_sparse_fill_empty_rows(
          sparse_indices_np: np.ndarray,
          sparse_values_np: np.ndarray,
          dense_shape_np: np.ndarray,
          default_value_np: np.ndarray,
      ) -> None:
          """
          This function verifies the relay output of sparse_fill_empty_rows with its expected output.
          """
          if use_dyn:
              sparse_indices = relay.var(
                  "sparse_indices",
                  shape=[relay.Any(), relay.Any()],
                  dtype=str(sparse_indices_np.dtype),
              )
              sparse_values = relay.var(
                  "sparse_values",
                  shape=[relay.Any()],
                  dtype=str(sparse_values_np.dtype),
              )
              dense_shape = relay.var(
                  "dense_shape",
                  shape=[relay.Any()],
                  dtype=str(dense_shape_np.dtype),
              )
              default_value = relay.var(
                  "default_value",
                  shape=[relay.Any()],
                  dtype=str(default_value_np.dtype),
              )
          else:
              sparse_indices = relay.var(
                  "sparse_indices",
                  relay.TensorType(sparse_indices_np.shape, str(sparse_indices_np.dtype)),
              )
              sparse_values = relay.var(
                  "sparse_values",
                  relay.TensorType(sparse_values_np.shape, str(sparse_values_np.dtype)),
              )
              dense_shape = relay.var(
                  "dense_shape",
                  relay.TensorType(dense_shape_np.shape, str(dense_shape_np.dtype)),
              )
              default_value = relay.var(
                  "default_value",
                  relay.TensorType(default_value_np.shape, str(default_value_np.dtype)),
              )
          z = relay.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value)
          func = relay.Function([sparse_indices, sparse_values, dense_shape, default_value], z)
          ref_res = ref_sparse_fill_empty_rows(
              sparse_indices_np,
              sparse_values_np,
              dense_shape_np,
              default_value_np,
          )
          (
              new_sparse_indices_infer_type,
              new_sparse_values_infer_type,
              empty_row_indicator_infer_type,
          ) = run_infer_type(z)
  
          assert new_sparse_indices_infer_type.checked_type.dtype == sparse_indices_np.dtype
          assert new_sparse_values_infer_type.checked_type.dtype == sparse_indices_np.dtype
          assert empty_row_indicator_infer_type.checked_type.dtype == "bool"
  
          verify_func(
              func,
              [sparse_indices_np, sparse_values_np, dense_shape_np, default_value_np],
              ref_res,
              [("llvm", tvm.cpu())],
          )
  
      verify_sparse_fill_empty_rows(
          sparse_indices.astype(dtype),
          sparse_values.astype(dtype),
          dense_shape.astype(dtype),
          default_value.astype(dtype),
      )
  
  
  if __name__ == "__main__":
      pytest.main([__file__])
  