// RUN: tf-opt -split-input-file -verify-diagnostics -tf-tpu-annotate-dynamic-shape-inputs %s | FileCheck %s

// Test that annotate the inputs of the cluster func to be dynamic shaped.

module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:worker/replica:0/task:0/device:TPU:0"]} {
   func.func @main(
      %arg0: tensor<2048xi64> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"},

};

TEST_F(PjRtExecutionUtilGpuTest, PreparePjRtExecutableArguments) {
  std::vector<const Tensor*> inputs;
  inputs.push_back(CreateDeviceTensor<int32_t>(TensorShape({1, 3}), {0, 0, 0}));
  inputs.push_back(CreateDeviceTensor<int32_t>(TensorShape({1, 3}), {1, 2, 3}));
  inputs.push_back(CreateDeviceTensor<int32_t>(TensorShape({1, 3}), {4, 5, 6}));
  int num_missing_prefix_ctx_inputs = 2;

	}
}

// TestSignV4TrimAll - tests the logic of TrimAll() function
func TestSignV4TrimAll(t *testing.T) {
	testCases := []struct {
		// Input.
		inputStr string
		// Expected result.
		result string
	}{
		{"本語", "本語"},
		{" abc ", "abc"},
		{" a b ", "a b"},
		{"a b ", "a b"},
		{"a  b", "a b"},
		{"a   b", "a b"},
		{"   a   b  c   ", "a b c"},

              Eq(TensorType_INT8));
  ASSERT_THAT(float_graph->tensors()->Get(float_op->inputs()->Get(1))->type(),
              Eq(TensorType_FLOAT32));
  EXPECT_THAT(subgraph->tensors[op->inputs[1]].get()->type,
              Eq(TensorType_INT8));

  // Verify FC bias should be int32 quantized.
  ASSERT_THAT(float_graph->tensors()->Get(float_op->inputs()->Get(2))->type(),

    // CHECK: %[[INPUT00:.*]] = "tf.IteratorGetNext"
    // CHECK-DAG: %[[SPACETODEPTH00:.*]] = "tf.SpaceToDepth"([[INPUT00:.*]]#0) <{block_size = 2 : i64, data_format = "NHWC"}> : (tensor<2x224x224x3xf32>) -> tensor<2x112x112x12xf32>

        value: "\n\003c:0\022\010c/Assign\032\010c/read:0"
      }
    }
  }
  signature_def {
    key: "classify_x2_to_y3"
    value {
      inputs {
        key: "inputs"
        value {
          name: "x2:0"
          dtype: DT_FLOAT
          tensor_shape {
            dim {
              size: -1
            }
            dim {

    }

    self._assertOpAndComposite([input_, filter_, bias],
                               tf.function(gen_mnist_ops.new_conv2d),
                               ops_defs._composite_conv_add_relu, kwargs)

  def test_new_conv2d_relu6(self):
    input_ = tf.random.uniform([1, 4, 4, 1])
    filter_ = tf.random.uniform([2, 2, 1, 8])
    bias = tf.zeros([8])
    kwargs = {
        'input_': input_,
        'filter_': filter_,

        private final ModelType<T> type;
        private final List<ModelReference<?>> inputs;

        private UnmanagedElementCreationAction(ModelRuleDescriptor descriptor, ModelReference<?> subject, List<ModelReference<?>> inputs, ModelType<T> type) {
            this.subject = subject;
            this.inputs = inputs;
            this.descriptor = descriptor;
            this.type = type;
        }

        @Override

            this.inputs = inputs;
            this.ruleSourceType = ruleSourceType;
            this.ruleExtractor = ruleExtractor;
        }

        @Override
        public ModelReference<?> getSubject() {
            return targetReference;
        }

        @Override
        public List<? extends ModelReference<?>> getInputs() {
            return inputs;
        }

        @Override

          absl::StrCat("tensor parsing error: ", alias));
    }

    inputs.emplace_back(feed_name, tensor);
  }

  if (!request_inputs.empty() &&
      seen_request_inputs.size() != request_inputs.size()) {
    return absl::InvalidArgumentError(absl::StrCat(
        "Inputs contains invalid name. Used request inputs: ",
        absl::StrJoin(seen_request_inputs, ","),