code: def _cg_check_shapes(a_mat, b_mat):
      if a_mat.shape[0] != a_mat.shape[1] or a_mat.shape.rank != 2:
          raise ValueError(
              "`a_mat` must be rank 2 square matrix; got shape {}.".format(
                  a_mat.shape
              )
          )
      if a_mat.shape[1] != b_mat.shape[0]:
          raise ValueError(
              "The dims of `a_mat` and `b_mat` are not compatible; "
              "got shapes {} and {}.".format(a_mat.shape, b_mat.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 1 for calling a global function with example usage: _cg_check_shapes(a_mat, b_mat) without return types