code: def mixed_mode_dot(a, b):
      """
      Computes the equivalent of `tf.einsum('ij,bjk->bik', a, b)`, but
      works for both dense and sparse inputs.
  
      :param a: Tensor or SparseTensor with rank 2.
      :param b: Tensor or SparseTensor with rank 3.
      :return: Tensor or SparseTensor with rank 3.
      """
      a_shp = tf.shape(a)
      b_shp = tf.shape(b)
  
      b_t = ops.transpose(b, (1, 2, 0))
      b_t = ops.reshape(b_t, tf.stack((b_shp[1], -1)))
      output = dot(a, b_t)
      output = ops.reshape(output, tf.stack((a_shp[0], b_shp[2], -1)))
      output = ops.transpose(output, (2, 0, 1))
  
      return output
  
file_path: spektral/spektral/layers/ops/matmul.py
imports: import tensorflow
  from tensorflow.keras import backend
  from tensorflow.python.ops.linalg.sparse import sparse
  from .None import ops
label_desc: Use image node_id 2 for calling a global function with example usage: mixed_mode_dot(a, b) and returns: output