Add an optional input signature for functions generated with defun.

An input signature is a possibly nested collection of `TensorSpec` objects declaring the shapes and dtypes of a function's arguments. tfe.defun propagates these shapes and dtypes to the graph function it generates for the traced Python function.
Since the shapes may be partially specified, this makes it possible to generate functions with partial shape information.

When an input signature is specified, every argument to the `defun`-wrapped function *must* be a Tensor. Input signatures cannot be specified for functions with keyword arguments, but positional arguments that are called in a keyword-like fashion are acceptable. For example, the following code snippet is valid

  @tfe.defun(input_signature=(TensorSpec([], tf.float32), TensorSpec([], tf.int32))
  def foo(b, a):
    ...

  foo(b=tf.constant(1.0), a=tf.constant(2))
  foo(a=tf.constant(2), b=tf.constant(1.0))

while the next code snippet is invalid

  @tfe.defun(input_signature=(TensorSpec([], tf.float32), TensorSpec([], tf.int32))
  def foo(b, **kwargs):
    ...

  # This will fail --- arbitrary kwargs are not allowed.
  foo(b=tf.constant(1.0), a=tf.constant(2))

This change also adds benchmarks that approximately measure the time taken to compute the cache key and verify the signature (and execute an empty function). Signatures introduce an overhead of ~100 us. The benchmarks for the non-signature path are the same as they were before this change.

entry {
  name: "MicroBenchmarks.benchmark_defun_with_signature"
  iters: 30000
  wall_time: 349.911403656
  extras {
    key: "examples_per_sec"
    value {
      double_value: 2857.86627572
    }
  }
}

entry {
  name: "MicroBenchmarks.benchmark_defun_with_signature_and_kwargs"
  iters: 30000
  wall_time: 360.46500206
  extras {
    key: "examples_per_sec"
    value {
      double_value: 2774.19442743
    }
  }
}

entry {
  name: "MicroBenchmarks.benchmark_defun_without_signature"
  iters: 30000
  wall_time: 259.087236722
  extras {
    key: "examples_per_sec"
    value {
      double_value: 3859.70383046
    }
  }
}

entry {
  name: "MicroBenchmarks.benchmark_defun_without_signature_and_with_kwargs"
  iters: 30000
  wall_time: 272.486400604
  extras {
    key: "examples_per_sec"
    value {
      double_value: 3669.90792121
    }
  }
}

PiperOrigin-RevId: 207617442