@test_util.run_deprecated_v1
  def testSquareBatch(self):
    self._testSquareBatch(dtypes_lib.bfloat16.as_numpy_dtype)
    self._testSquareBatch(np.float32)
    self._testSquareBatch(np.float64)