code: def _check_shapes(a_inv: types.Tensor, u: types.Tensor):
      if a_inv.shape[0] != a_inv.shape[1] or a_inv.shape.rank != 2:
          raise ValueError(
              "`a_inv` must be rank 2 square matrix; got shape {}.".format(
                  a_inv.shape
              )
          )
      if u.shape.rank != 2:
          raise ValueError(
              "`u` must be rank 2 matrix; got shape {}.".format(u.shape)
          )
      if a_inv.shape[1] != u.shape[1]:
          raise ValueError(
              "`a_inv` and `u` must have shapes [m, m] and [n, m]; "
              "got shapes {} and {}.".format(a_inv.shape, u.shape)
          )
  
file_path: agents/tf_agents/bandits/policies/linalg.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 3 for calling a global function with example usage: _check_shapes(a_inv, u) without return types