Emitting MLIR of training step function

[Laszlo Kindrat]

Hi everyone,

I'm trying to convert a tf.function wrapping an entire training step into MLIR, like in this snippet. I'm using tf 2.5.

The resulting MLIR has the correct shape and type information for the inputs that were explicitly specified (x and y), but the inputs for the model variables all have types tensor<*x!tf.resource>. I wonder if there is a way to keep the shape/type info for these inputs in the emitted MLIR.

Thanks in advance!