message: Add tf.BatchFunction to Tensorflow MLIR ODS.
diff: @@ -1256,4 +1256,74 @@ def TF_FusedBatchNormV3Op : TF_FusedBatchNormOpBase<"FusedBatchNormV3"> {
  );
  }
  +def TF_BatchFunctionOp : TF_Op<"BatchFunction", [AttrSizedOperandSegments]> {
  + let summary = [{
  +Batches all the inputs tensors to the computation done by the function.
  + }];
  +
  + let description = [{
  +So, for example, in the following code
  +
  + ```python
  +
  + # This input will be captured.
  + y = tf.placeholder_with_default(1.0, shape=[])
  +
  + @tf.Defun(tf.float32)
  + def computation(a):
  + return tf.matmul(a, a) + y
  +
  + b = gen_batch_ops.batch_function(
  + f=computation
  + in_tensors=[a],
  + captured_tensors=computation.captured_inputs,
  + Tout=[o.type for o in computation.definition.signature.output_arg],
  + num_batch_threads=1,
  + max_batch_size=10,
  + batch_timeout_micros=100000, # 100ms
  + allowed_batch_sizes=[3, 10],
  + batching_queue="")
  +
  +If more than one session.run call is simultaneously trying to compute `b`
  +the values of `a` will be gathered, non-deterministically concatenated
  +along the first axis, and only one thread will run the computation.
  +
  +Assumes that all arguments of the function are Tensors which will be batched
  +along their first dimension.
  +
  +Arguments that are captured, are not batched. The session.run call which does
  +the concatenation, will use the values of the captured tensors available to it.
  +Therefore, typical uses of captured tensors should involve values which remain
  +unchanged across session.run calls. Inference is a good example of this.
  +
  +SparseTensor is not supported. The return value of the decorated function
  +must be a Tensor or a list/tuple of Tensors.
  + }];
  +
  + let arguments = (ins
  + Variadic<TF_Tensor>:$in_tensors,
  + Variadic<TF_Tensor>:$captured_tensors,
  +
  + SymbolRefAttr:$f,
  + I64Attr:$num_batch_threads,
  + I64Attr:$max_batch_size,
  + I64Attr:$batch_timeout_micros,
  + DefaultValuedAttr<I64Attr, "10">:$max_enqueued_batches,
  + DefaultValuedAttr<I64ArrayAttr, "{}">:$allowed_batch_sizes,
  + StrAttr:$container,
  + StrAttr:$shared_name,
  + StrAttr:$batching_queue,
  + DefaultValuedAttr<BoolAttr, "false">:$enable_large_batch_splitting,
  + I32ElementsAttr:$operand_segment_sizes
  + );
  +
  + let results = (outs
  + Variadic<TF_Tensor>:$out_tensors
  + );
  +
  + TF_DerivedOperandTypeListAttr Tin = TF_DerivedOperandTypeListAttr<0>;
  + TF_DerivedOperandTypeListAttr Tcaptured = TF_DerivedOperandTypeListAttr<1>;
  + TF_DerivedResultTypeListAttr Tout = TF_DerivedResultTypeListAttr<0>;
  +}
  +
  #endif // TF_OPS
  