Address several autograph failed issues for TF2

[WeichenXu123]

Currently, for some function, TF2 autograph will fail. See https://github.com/tensorflow/tensorflow/issues/35765 https://github.com/tensorflow/tensorflow/issues/30149 https://github.com/tensorflow/autograph/issues/3

If autograph failed, the functions will be run eagerly and TF cannot optimize them. So we'd better address them.

Manually test

df1 = spark.range(100)
from petastorm.spark import make_spark_converter

# Set a cache directory on DBFS FUSE for intermediate data.
spark.conf.set("petastorm.spark.converter.parentCacheDirUrl", "file:///dbfs/ml/tmp/petastorm/QA/bugs/")
converter1 = make_spark_converter(df1)

with converter1.make_tf_dataset(num_epochs=1) as dataset:
  for batch in dataset:
    print(batch.id)

• Before Output includes:

  
  WARNING:tensorflow:AutoGraph could not transform <function _NamedtupleCache.get at 0x7f0bfbe6f200> and will run it as-is.
  Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.
  Cause: expected exactly one node node, found [<gast.gast.FunctionDef object at 0x7f0bfad5d050>, <gast.gast.Return object at 0x7f0bfad5d7d0>]
  WARNING:tensorflow:AutoGraph could not transform <function make_petastorm_dataset.<locals>.<lambda> at 0x7f0bf8da03b0> and will run it as-is.
  Cause: could not parse the source code:
  
          .map(lambda row: _set_shape_to_named_tuple(reader.schema, row, reader.batched_output))

This error may be avoided by creating the lambda in a standalone statement.

* After
The warnings listed above disappear.

[selitvin]

How did you find these failures? Just by running in your external environment? If so, is it hard to add a test to make sure we don't break autograph going forward?

[codecov[bot]]

Codecov Report

Merging #542 into master will increase coverage by 0.00%. The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master     #542   +/-   ##
=======================================
  Coverage   86.15%   86.16%           
=======================================
  Files          87       87           
  Lines        4932     4935    +3     
  Branches      787      786    -1     
=======================================
+ Hits         4249     4252    +3     
  Misses        556      556           
  Partials      127      127           

Impacted Files	Coverage Δ
petastorm/tf_utils.py	88.65% <100.00%> (+0.24%)	:arrow_up:
petastorm/unischema.py	95.79% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more Δ = absolute <relative> (impact), ø = not affected, ? = missing data Powered by Codecov. Last update f5d6ea1...5b53213. Read the comment docs.