code: def test_vecmat():
      @triton.jit
      def batched_vecmat(
          # inputs
          A,  # shape: [dim_m, dim_k]
          B,  # shape: [dim_m, dim_n, dim_k]
          # dimensions
          dim_m,
          dim_n,
          dim_k,
          # outputs
          output,
          # block information
          block_m: tl.constexpr,
          block_n: tl.constexpr,
          block_k: tl.constexpr,
      ):
          m_index = tl.program_id(0)
          n_index = tl.program_id(1)
          # Output tile
          output_tile = (m_index * block_m + tl.arange(0, block_m))[
              :, None
          ] * dim_n + (n_index * block_n + tl.arange(0, block_n))[
              None, :
          ]
  
          vecmat = tl.zeros(
              [block_m, block_n], dtype=A.dtype.element_ty
          )
          k_blocks = dim_k // block_k
          for k_index in range(k_blocks):
              # Load A tile
              a_tile = (m_index * block_m + tl.arange(0, block_m))[
                  :, None
              ] * dim_k + (k_index * block_k + tl.arange(0, block_k))[
                  None, :
              ]
              a = tl.load(A + a_tile)
  
              # Load B tile, transposed to [n, m, k] in order to broadcast A on a
              # leading dimension.
              b_tile = (
                  (m_index * block_m + tl.arange(0, block_m))[
                      None, :, None
                  ]
                  * dim_n
                  * dim_k
                  + (n_index * block_n + tl.arange(0, block_n))[
                      :, None, None
                  ]
                  * dim_k
                  + (k_index * block_k + tl.arange(0, block_k))[
                      None, None, :
                  ]
              )
              b = tl.load(B + b_tile)
  
              expanded_a, _ = tl.broadcast(a, b)
              vecmat += tl.trans(tl.sum(expanded_a * b, axis=2))
  
          tl.store(output + output_tile, vecmat)
  
      M, N, K = 128, 128, 128
      block_m, block_n, block_k = 16, 32, 64
  
      rs = RandomState(17)
      A_vec = rs.randint(0, 4, (M, K)).astype("float32")
      B_vec = rs.randint(0, 4, (M, N, K)).astype("float32")
      A = A_vec
      B = B_vec
  
      A_tri = torch.tensor(A, device="cuda")
      B_tri = torch.tensor(B, device="cuda")
      C_tri = torch.zeros((M, N), dtype=torch.float32, device="cuda")
  
      grid = (M // block_m, N // block_n)
  
      batched_vecmat[grid](
          A_tri,
          B_tri,
          M,
          N,
          K,
          C_tri,  #
          block_m=block_m,
          block_n=block_n,
          block_k=block_k,  #
          num_warps=4,
          num_stages=1,
      )
  
      A_expanded = A[:, np.newaxis, :]
      A_broadcasted = np.broadcast_to(A_expanded, (M, N, K))
      AB = A_broadcasted * B
      C_ref = np.sum(AB, axis=2)
  
      np.testing.assert_allclose(
          C_ref, C_tri.cpu().numpy(), rtol=0.01, atol=1e-3
      )
  
file_path: triton/python/test/regression/test_functional_regressions.py
imports: import numpy
  import pytest
  import torch
  from numpy.random import RandomState
  import triton
  import triton.language
label_desc: Use image node_id 2 for calling a global function with example usage: test_vecmat() without return types