code: def dot(a, b):
      """
      Computes a @ b, for a, b of the same rank (both 2 or both 3).
  
      If the rank is 2, then the innermost dimension of `a` must match the
      outermost dimension of `b`.
      If the rank is 3, the first dimension of `a` and `b` must be equal and the
      function computes a batch matmul.
  
      Supports both dense and sparse multiplication (including sparse-sparse).
  
      :param a: Tensor or SparseTensor with rank 2 or 3.
      :param b: Tensor or SparseTensor with same rank as b.
      :return: Tensor or SparseTensor with rank 2 or 3.
      """
      a_ndim = K.ndim(a)
      b_ndim = K.ndim(b)
      assert (
          a_ndim == b_ndim
      ), "Expected equal ranks, got {} and {}" "".format(a_ndim, b_ndim)
      a_is_sparse = K.is_sparse(a)
      b_is_sparse = K.is_sparse(b)
  
      # Handle cases: rank 2 sparse-dense, rank 2 dense-sparse
      # In these cases we can use the faster sparse-dense matmul of tf.sparse
      if a_ndim == 2:
          if a_is_sparse and not b_is_sparse:
              return tf.sparse.sparse_dense_matmul(a, b)
          if not a_is_sparse and b_is_sparse:
              return ops.transpose(
                  tf.sparse.sparse_dense_matmul(
                      ops.transpose(b), ops.transpose(a)
                  )
              )
  
      # Handle cases: rank 2 sparse-sparse, rank 3 sparse-dense,
      # rank 3 dense-sparse, rank 3 sparse-sparse
      # In these cases we can use the tfsp.CSRSparseMatrix implementation (slower,
      # but saves memory)
      if a_is_sparse:
          a = tfsp.CSRSparseMatrix(a)
      if b_is_sparse:
          b = tfsp.CSRSparseMatrix(b)
      if a_is_sparse or b_is_sparse:
          out = tfsp.matmul(a, b)
          if hasattr(out, "to_sparse_tensor"):
              return out.to_sparse_tensor()
          else:
              return out
  
      # Handle case: rank 2 dense-dense, rank 3 dense-dense
      # Here we use the standard dense operation
      return tf.matmul(a, b)
  
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 1 for calling a global function with example usage: dot(a, b) and returns: tf, tf, ops, out, out