Shape related code updated to address TODO(b/208879020)#270

pritamdodeja wants to merge 3 commits intotensorflow:mastertensorflow/transform:masterfrom

pritamdodeja:shape_fixpritamdodeja/transform:shape_fixCopy head branch name to clipboard