I’m not sure how the TDD layer works in TensorFlow, but as given in the code snippet, the linear layer will be applied to all samples in dim1 separately.
If that’s what TDD does, then it should be equal. A verification with some constant inputs would be nice, in case you have a TF installation